Created
August 3, 2024 16:37
-
-
Save RandyGaul/b58316c2d3ffcdf680cf6c4c6572698b to your computer and use it in GitHub Desktop.
priority queue implementation in lua
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| -- Piorirty queue, used for implementing other algorithms such as | |
| -- prioritized message queues, or the A* algorithm. | |
| -- | |
| -- Create a new priority q like so: | |
| -- | |
| -- q = prio_q() | |
| -- | |
| -- Add elements to the queue, each has a cost associated. THe cost is | |
| -- used to sort the elements, where the lowest cost is considered the | |
| -- highest priority (first out when pop() is called). | |
| -- | |
| -- q:add("A", 5) | |
| -- q:add("B", 10) | |
| -- q:add("C", 3) | |
| -- | |
| -- Pop elements from the queue to fetch them in sorted order. | |
| -- | |
| -- v, cost = q:pop() | |
| -- print(v, cost) -- prints "C", 3 | |
| prio_q = class:extend() | |
| function prio_q:new() | |
| self.values = {} | |
| self.costs = {} | |
| end | |
| -- Push a value onto the queue with a cost. | |
| -- ...lower the cost means it will be pop'd at higher priority. | |
| function prio_q:push(val, cost) | |
| local values, costs = self.values, self.costs | |
| local i = #values + 1 | |
| values[i] = val | |
| costs[i] = cost | |
| -- Heapify up. | |
| while i > 1 do | |
| local parent = math.floor(i / 2) | |
| if self:predicate_min(i, parent) < 0 then | |
| self:swap(i, parent) | |
| i = parent | |
| else | |
| break | |
| end | |
| end | |
| end | |
| -- Pop the value off of the queue with the minimum cost. | |
| -- ...return v, cost | |
| function prio_q:pop() | |
| local values, costs = self.values, self.costs | |
| local count = #values | |
| if count == 0 then return nil, nil end | |
| local val = values[1] | |
| local cost = costs[1] | |
| -- Move the last element to the root. | |
| values[1] = values[count] | |
| costs[1] = costs[count] | |
| -- Remove last element. | |
| values[count] = nil | |
| costs[count] = nil | |
| -- Heapify down. | |
| local i = 2 | |
| while true do | |
| local left = 2 * i | |
| local right = 2 * i + 1 | |
| local smallest = i | |
| if left <= #values and self:predicate_min(left, smallest) < 0 then | |
| smallest = left | |
| end | |
| if right <= #values and self:predicate_min(right, smallest) < 0 then | |
| smallest = right | |
| end | |
| if smallest ~= i then | |
| self:swap(i, smallest) | |
| i = smallest | |
| else | |
| break | |
| end | |
| end | |
| return val, cost | |
| end | |
| -- Return the cnumber of elements in the queue. | |
| function prio_q:count() | |
| return #self.values | |
| end | |
| -- Clear the queue. | |
| function prio_q:clear() | |
| self.values = {} | |
| self.costs = {} | |
| end | |
| -- This is an internal function used for sorting. | |
| function prio_q:predicate_min(iA, iB) | |
| local costs = self.costs | |
| local costA = costs[iA] | |
| local costB = costs[iB] | |
| return costA < costB and -1 or (costA > costB and 1 or 0) | |
| end | |
| -- This is an internal function used for sorting. | |
| function prio_q:swap(iA, iB) | |
| local values, costs = self.values, self.costs | |
| values[iA], values[iB] = values[iB], values[iA] | |
| costs[iA], costs[iB] = costs[iB], costs[iA] | |
| end | |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment