Merge pull request #8011 from JuliaLang/teh/pq

PriorityQueue: use O(logN) algorithm for dequeue!(pq, key)

base/collections.jl

@@ -1,10 +1,11 @@
 
 module Collections
 
-import Base: setindex!, done, get, haskey, isempty, length, next, getindex, start
+import Base: setindex!, done, get, hash, haskey, isempty, length, next, getindex, start
 import ..Order: Forward, Ordering, lt
 
 export
+    Pair,
     PriorityQueue,
     dequeue!,
     enqueue!,
@@ -106,37 +107,48 @@ end
 # PriorityQueue
 # -------------
 
+immutable Pair{A,B}
+    a::A
+    b::B
+end
+
+function hash(p::Pair, h::Uint)
+    h = hash(p.a, h)
+    hash(p.b, h)
+end
+hash(p::Pair) = hash(p, zero(Uint))
+
 # A PriorityQueue that acts like a Dict, mapping values to their priorities,
 # with the addition of a dequeue! function to remove the lowest priority
 # element.
-type PriorityQueue{K,V} <: Associative{K,V}
+type PriorityQueue{K,V,O<:Ordering} <: Associative{K,V}
     # Binary heap of (element, priority) pairs.
-    xs::Array{(K, V), 1}
-    o::Ordering
+    xs::Array{Pair{K,V}, 1}
+    o::O
 
     # Map elements to their index in xs
     index::Dict{K, Int}
 
-    function PriorityQueue(o::Ordering)
-        new(Array((K, V), 0), o, Dict{K, Int}())
+    function PriorityQueue(o::O)
+        new(Array(Pair{K,V}, 0), o, Dict{K, Int}())
     end
 
-    PriorityQueue() = PriorityQueue{K,V}(Forward)
+    PriorityQueue() = PriorityQueue{K,V,O}(Forward)
 
     function PriorityQueue(ks::AbstractArray{K}, vs::AbstractArray{V},
-                           o::Ordering)
+                           o::O)
         # TODO: maybe deprecate
         if length(ks) != length(vs)
             error("key and value arrays must have equal lengths")
         end
-        PriorityQueue{K,V}(zip(ks, vs), o)
+        PriorityQueue{K,V,O}(zip(ks, vs), o)
     end
 
-    function PriorityQueue(itr, o::Ordering)
-        xs = Array((K, V), length(itr))
+    function PriorityQueue(itr, o::O)
+        xs = Array(Pair{K,V}, length(itr))
         index = Dict{K, Int}()
         for (i, (k, v)) in enumerate(itr)
-            xs[i] = (k, v)
+            xs[i] = Pair{K,V}(k, v)
             if haskey(index, k)
                 error("PriorityQueue keys must be unique")
             end
@@ -153,36 +165,37 @@ type PriorityQueue{K,V} <: Associative{K,V}
     end
 end
 
-PriorityQueue(o::Ordering=Forward) = PriorityQueue{Any,Any}(o)
+PriorityQueue(o::Ordering=Forward) = PriorityQueue{Any,Any,typeof(o)}(o)
+PriorityQueue{K,V}(::Type{K}, ::Type{V}, o::Ordering=Forward) = PriorityQueue{K,V,typeof(o)}(o)
 
 # TODO: maybe deprecate
 PriorityQueue{K,V}(ks::AbstractArray{K}, vs::AbstractArray{V},
-                   o::Ordering=Forward) = PriorityQueue{K,V}(ks, vs, o)
+                   o::Ordering=Forward) = PriorityQueue{K,V,typeof(o)}(ks, vs, o)
 
-PriorityQueue{K,V}(kvs::Associative{K,V}, o::Ordering=Forward) = PriorityQueue{K,V}(kvs, o)
+PriorityQueue{K,V}(kvs::Associative{K,V}, o::Ordering=Forward) = PriorityQueue{K,V,typeof(o)}(kvs, o)
 
-PriorityQueue{K,V}(a::AbstractArray{(K,V)}, o::Ordering=Forward) = PriorityQueue{K,V}(a, o)
+PriorityQueue{K,V}(a::AbstractArray{(K,V)}, o::Ordering=Forward) = PriorityQueue{K,V,typeof(o)}(a, o)
 
 length(pq::PriorityQueue) = length(pq.xs)
 isempty(pq::PriorityQueue) = isempty(pq.xs)
 haskey(pq::PriorityQueue, key) = haskey(pq.index, key)
-peek(pq::PriorityQueue) = pq.xs[1]
+peek(pq::PriorityQueue) = (kv = pq.xs[1]; (kv.a, kv.b))
 
 
 function percolate_down!(pq::PriorityQueue, i::Integer)
     x = pq.xs[i]
     @inbounds while (l = heapleft(i)) <= length(pq)
         r = heapright(i)
-        j = r > length(pq) || lt(pq.o, pq.xs[l][2], pq.xs[r][2]) ? l : r
-        if lt(pq.o, pq.xs[j][2], x[2])
-            pq.index[pq.xs[j][1]] = i
+        j = r > length(pq) || lt(pq.o, pq.xs[l].b, pq.xs[r].b) ? l : r
+        if lt(pq.o, pq.xs[j].b, x.b)
+            pq.index[pq.xs[j].a] = i
             pq.xs[i] = pq.xs[j]
             i = j
         else
             break
         end
     end
-    pq.index[x[1]] = i
+    pq.index[x.a] = i
     pq.xs[i] = x
 end
 
@@ -191,36 +204,48 @@ function percolate_up!(pq::PriorityQueue, i::Integer)
     x = pq.xs[i]
     @inbounds while i > 1
         j = heapparent(i)
-        if lt(pq.o, x[2], pq.xs[j][2])
-            pq.index[pq.xs[j][1]] = i
+        if lt(pq.o, x.b, pq.xs[j].b)
+            pq.index[pq.xs[j].a] = i
             pq.xs[i] = pq.xs[j]
             i = j
         else
             break
         end
     end
-    pq.index[x[1]] = i
+    pq.index[x.a] = i
     pq.xs[i] = x
 end
 
+# Equivalent to percolate_up! with an element having lower priority than any other
+function force_up!(pq::PriorityQueue, i::Integer)
+    x = pq.xs[i]
+    @inbounds while i > 1
+        j = heapparent(i)
+        pq.index[pq.xs[j].a] = i
+        pq.xs[i] = pq.xs[j]
+        i = j
+    end
+    pq.index[x.a] = i
+    pq.xs[i] = x
+end
 
 function getindex{K,V}(pq::PriorityQueue{K,V}, key)
-    pq.xs[pq.index[key]][2]
+    pq.xs[pq.index[key]].b
 end
 
 
 function get{K,V}(pq::PriorityQueue{K,V}, key, deflt)
     i = get(pq.index, key, 0)
-    i == 0 ? deflt : pq.xs[i][2]
+    i == 0 ? deflt : pq.xs[i].b
 end
 
 
 # Change the priority of an existing element, or equeue it if it isn't present.
 function setindex!{K,V}(pq::PriorityQueue{K, V}, value, key)
     if haskey(pq, key)
         i = pq.index[key]
-        _, oldvalue = pq.xs[i]
-        pq.xs[i] = (key, value)
+        oldvalue = pq.xs[i].b
+        pq.xs[i] = Pair{K,V}(key, value)
         if lt(pq.o, oldvalue, value)
             percolate_down!(pq, i)
         else
@@ -229,6 +254,7 @@ function setindex!{K,V}(pq::PriorityQueue{K, V}, value, key)
     else
         enqueue!(pq, key, value)
     end
+    value
 end
 
 
@@ -237,7 +263,7 @@ function enqueue!{K,V}(pq::PriorityQueue{K,V}, key, value)
         error("PriorityQueue keys must be unique")
     end
 
-    push!(pq.xs, (key, value))
+    push!(pq.xs, Pair{K,V}(key, value))
     pq.index[key] = length(pq)
     percolate_up!(pq, length(pq))
     pq
@@ -249,19 +275,17 @@ function dequeue!(pq::PriorityQueue)
     y = pop!(pq.xs)
     if !isempty(pq)
         pq.xs[1] = y
-        pq.index[pq.xs[1][1]] = 1
+        pq.index[y.a] = 1
         percolate_down!(pq, 1)
     end
-    delete!(pq.index, x[1])
-    x[1]
+    delete!(pq.index, x.a)
+    x.a
 end
 
 function dequeue!(pq::PriorityQueue, key)
-    idx = pop!(pq.index, key)  # throws key error if missing
-    deleteat!(pq.xs, idx)
-    for (k,v) in pq.index
-        (v >= idx) && (pq.index[k] = (v-1))
-    end
+    idx = pq.index[key]
+    force_up!(pq, idx)
+    dequeue!(pq)
     key
 end
 
@@ -273,7 +297,7 @@ done(pq::PriorityQueue, i) = done(pq.index, i)
 
 function next(pq::PriorityQueue, i)
     (k, idx), i = next(pq.index, i)
-    return ((k, pq.xs[idx][2]), i)
+    return ((k, pq.xs[idx].b), i)
 end
 
 

test/priorityqueue.jl

@@ -12,14 +12,26 @@ function test_issorted!(pq::PriorityQueue, priorities)
     end
 end
 
+function test_isrequested!(pq::PriorityQueue, keys)
+    i = 0
+    while !isempty(pq)
+        krqst =  keys[i+=1]
+        krcvd = dequeue!(pq, krqst)
+        @test krcvd == krqst
+    end
+end
+
 pmax = 1000
 n = 10000
-priorities = Dict(1:n, rand(1:pmax, n))
+r = rand(1:pmax, n)
+priorities = Dict(1:n, r)
 
 # building from a dict
 pq = PriorityQueue(priorities)
 test_issorted!(pq, priorities)
 
+pq = PriorityQueue(priorities)
+test_isrequested!(pq, 1:n)
 
 # building from two lists
 ks, vs = 1:n, rand(1:pmax, n)