Merge pull request #269 from CarloLucibello/heap

indexing interface for MutableBinaryHeap
JuliaCollections · Jul 13, 2017 · e8c2e9b · e8c2e9b
2 parents 7f5032d + 74a5f2c
commit e8c2e9b
Show file tree

Hide file tree

Showing 3 changed files with 15 additions and 1 deletion.
diff --git a/src/heaps/mutable_binary_heap.jl b/src/heaps/mutable_binary_heap.jl
@@ -226,6 +226,7 @@ pop!{T}(h::MutableBinaryHeap{T}) = _binary_heap_pop!(h.comparer, h.nodes, h.node
     update!{T}(h::MutableBinaryHeap{T}, i::Int, v::T)
 
 Replace the element at index `i` in heap `h` with `v`.
+This is equivalent to `h[i]=v`.
 """
 function update!{T}(h::MutableBinaryHeap{T}, i::Int, v::T)
     nodes = h.nodes
@@ -241,3 +242,6 @@ function update!{T}(h::MutableBinaryHeap{T}, i::Int, v::T)
         _heap_bubble_down!(comp, nodes, nodemap, nd_id)
     end
 end
+
+setindex!(h::MutableBinaryHeap, v, i::Int) = update!(h, i, v)
+getindex(h::MutableBinaryHeap, i::Int) = h.nodes[h.node_map[i]].value
diff --git a/test/runtests.jl b/test/runtests.jl
@@ -31,7 +31,8 @@ tests = ["int_set",
          "multi_dict",
          "circular_buffer",
          "sorting",
-         "priorityqueue"]
+         "priorityqueue"
+        ]
 
 if length(ARGS) > 0
     tests = ARGS

diff --git a/test/test_mutable_binheap.jl b/test/test_mutable_binheap.jl
@@ -192,5 +192,14 @@ for (hf,m) = [(mutable_binary_minheap,-2.0), (mutable_binary_maxheap,2.0)]
         @test verify_heap(h)
         @test isequal(list_values(h), xs)
         @test top_with_handle(h) == (v, i)
+
+        i = rand(1:100)
+        v = rand()
+        h[i] = v
+        xs[i] = v
+        @test length(h) == 100
+        @test verify_heap(h)
+        @test isequal(list_values(h), xs)
+        @test v == h[i]
     end
 end