Indexed heap with decrease-key in Python

Why this matters

A route planner finds the shortest way from a start city to every other city. It keeps a list of cities with a best-known distance for each one. Every time it finds a shorter road into a city already on the list, that city's distance goes down. The city has to move up in the queue, so the planner can take the closest city next.

That step is called decrease-key. A key is the number a heap uses to order its items. In this lesson the key is a priority, and a smaller number means more urgent. Decrease-key means: an item already in the heap gets a smaller key.

Python's heapq module keeps a heap in an ordinary list, but it has no decrease-key operation. To lower one item's key with heapq, you push a second copy of the item with the new key and skip the old copy when it comes out. This works, and the mistakes section describes it. The cost is that the heap can hold many old copies. An indexed heap keeps exactly one copy of each item. It gives each item a whole-number id, from 0 up to n minus 1, and it keeps a second list that records where each id sits in the heap. With that second list, the program jumps straight to the item and moves it up. The cost is O(log n), which means the number of steps grows like the number of times you can halve n. For a thousand items, that is about 10 steps.

This lesson builds the indexed heap in Python, traces it by hand, and uses it to run Dijkstra's shortest-path method. It does not re-teach how a heap is stored in a list. The lesson on the heap property and list layout covers that.

Start with a shelf of labelled boxes

Picture a shelf with boxes. Each box has a label, such as A, and a number written on it. The smallest number must be at the front of the shelf, and the rest is only partly ordered, as a heap keeps it.

Now someone says: "Box C's number is now 1." You need to find box C first. If you have to look at every box to find it, that search costs time. The heap itself is fast at moving a box once you hold it. The slow part is finding the box.

An indexed heap solves this with a sticky note on each box that says its current slot on the shelf. When you need box C, you read its note and go straight to that slot.

Three lists that agree

The indexed heap keeps three lists. Each item has an id from 0 to n minus 1, so the lists can be indexed by id.

key[id] holds the current priority of the item with that id.

heap[p] holds the id of the item sitting at position p in the heap. Position 0 is the top.

pos[id] holds the position of the item with that id. When the item is not in the heap, the value is -1.

The three lists always agree. If heap[3] is 7, then pos[7] is 3. Every time two items swap places, the code must update pos for both of them. That rule is the whole trick.

A worked trace by hand

Start with four items named A, B, C and D, with ids 0, 1, 2 and 3. Insert them with these keys: A is 5, B is 3, C is 8, and D is 6. Then lower C to 1.

Insert A with key 5. The heap has one item, so A sits at position 0. The heap is [A5].

Insert B with key 3. B goes to the end at position 1. Its parent is at position 0, which holds A with key 5. Since 3 is smaller than 5, B and A swap. Now the heap is [B3, A5].

Insert C with key 8. C goes to position 2. Its parent is B with key 3. Since 8 is not smaller than 3, nothing moves. The heap is [B3, A5, C8].

Insert D with key 6. D goes to position 3. Its parent is at position 1, which holds A with key 5. Since 6 is not smaller than 5, nothing moves. The heap is [B3, A5, C8, D6].

Now lower C from 8 to 1. C is at position 2. Its parent at position 0 holds B with key 3. Since 1 is smaller than 3, C and B swap. Now C is at position 0. The heap is [C1, A5, B3, D6].

Only one swap happened. The code did not search the heap for C, because pos[C] already said position 2.

The Python class

The class below stores ids and keys in lists. It has the methods you need: insert, poll (remove the smallest), decrease-key, and a few checks. The dump method returns the heap as text, so you can see the trace. The file ends with a driver that runs the trace above.

Python · editable

import math

class IndexedMinHeap:
    def __init__(self, capacity):
        self.heap = [0] * capacity  # heap[p] = id of the item at position p
        self.pos = [-1] * capacity  # pos[id] = position of item id, or -1 if not in the heap
        self.key = [0] * capacity  # key[id] = current priority of item id
        self.size = 0

def __len__(self):
        return self.size

def contains(self, item_id):
        return self.pos[item_id] != -1

def priority_of(self, item_id):
        return self.key[item_id]

def insert(self, item_id, priority):
        if self.contains(item_id):
            raise ValueError(f"already in heap: {item_id}")
        self.key[item_id] = priority
        self.heap[self.size] = item_id
        self.pos[item_id] = self.size
        self.size += 1
        self._sift_up(self.pos[item_id])

def peek_id(self):
        if self.size == 0:
            raise IndexError("heap is empty")
        return self.heap[0]

def poll_id(self):
        if self.size == 0:
            raise IndexError("heap is empty")
        top = self.heap[0]
        self._swap(0, self.size - 1)
        self.size -= 1
        self.pos[top] = -1
        if self.size > 0:
            self._sift_down(0)
        return top

def decrease_key(self, item_id, new_priority):
        if not self.contains(item_id):
            raise ValueError(f"not in heap: {item_id}")
        if new_priority > self.key[item_id]:
            raise ValueError(f"new priority {new_priority} is larger than {self.key[item_id]}")
        self.key[item_id] = new_priority
        self._sift_up(self.pos[item_id])

def dump(self):
        parts = [f"{item_id}:{self.key[item_id]}" for item_id in self.heap[: self.size]]
        return "[" + ", ".join(parts) + "]"

def _sift_up(self, p):
        while p > 0:
            parent = (p - 1) // 2
            if self.key[self.heap[parent]] <= self.key[self.heap[p]]:
                break
            self._swap(p, parent)
            p = parent

def _sift_down(self, p):
        while True:
            left = 2 * p + 1
            right = left + 1
            smallest = p
            if left < self.size and self.key[self.heap[left]] < self.key[self.heap[smallest]]:
                smallest = left
            if right < self.size and self.key[self.heap[right]] < self.key[self.heap[smallest]]:
                smallest = right
            if smallest == p:
                break
            self._swap(p, smallest)
            p = smallest

def _swap(self, a, b):
        id_a = self.heap[a]
        id_b = self.heap[b]
        self.heap[a] = id_b
        self.heap[b] = id_a
        self.pos[id_b] = a
        self.pos[id_a] = b

if __name__ == "__main__":
    h = IndexedMinHeap(4)
    h.insert(0, 5)
    print("after insert A=5:", h.dump())
    h.insert(1, 3)
    print("after insert B=3:", h.dump())
    h.insert(2, 8)
    print("after insert C=8:", h.dump())
    h.insert(3, 6)
    print("after insert D=6:", h.dump())
    h.decrease_key(2, 1)
    print("after C to 1:", h.dump())
    print("C still in heap:", h.contains(2), "priority", h.priority_of(2))
    print("poll:", h.poll_id())
    print("after poll:", h.dump())
    print("poll:", h.poll_id())
    print("after poll:", h.dump())

Save the file as indexed_heap.py and run python3 indexed_heap.py. The output is:

after insert A=5: [0:5]
after insert B=3: [1:3, 0:5]
after insert C=8: [1:3, 0:5, 2:8]
after insert D=6: [1:3, 0:5, 2:8, 3:6]
after C to 1: [2:1, 0:5, 1:3, 3:6]
C still in heap: True priority 1
poll: 2
after poll: [1:3, 0:5, 3:6]
poll: 1
after poll: [0:5, 3:6]

The output matches the hand trace. Each entry reads id:key. The id is 0 for A, 1 for B, 2 for C and 3 for D.

Polling then removes the top, which is C. The heap moves its last item to the top and sinks it down.

Reading the class

The lists are created with [0] * capacity. That makes a list of capacity zeros, which is fixed in size. The pos list starts with -1 everywhere, meaning no item is in the heap yet.

The parent of position p is (p - 1) // 2, using floor division. The children of p are 2 * p + 1 and 2 * p + 2.

_swap is the only method that moves items, and it updates pos for both items. If you move items anywhere else, you must update pos yourself.

decrease_key calls _sift_up (swap the item upward while its parent is larger), never _sift_down. Lowering a key can only break the rule with the parent above. The children were already bigger than the old key, so they are still bigger than the new one.

dump uses a slice, self.heap[: self.size], to read only the part of the list that holds the heap. The list itself is longer than the heap when the capacity is larger than the number of items.

A larger example

Now take seven items, with ids 0 to 6 and keys 50, 40, 30, 20, 10, 60 and 70. This heap has three levels: position 0 at the top, positions 1 and 2 below it, and positions 3 to 6 at the bottom. The inserts run the same way as before. Each new item goes to the end, then swaps with its parent while it is smaller. The output after all seven inserts is:

after inserts: [4:10, 3:20, 1:40, 0:50, 2:30, 5:60, 6:70]

Item 6 has key 70 and sits at position 6. Its parent is position 2, which holds item 1 with key 40. Lower item 6 to key 5. The code reads pos[6], which is 6, and goes to the parent at position 2. Five is smaller than 40, so they swap. Item 6 is now at position 2, and its parent at position 0 holds item 4 with key 10. Five is smaller, so they swap again. Item 6 reaches the top after two swaps. Each level it climbs costs one swap, so a decrease-key costs at most one swap per level.

after id 6 to 5: [6:5, 3:20, 4:10, 0:50, 2:30, 5:60, 1:40]
root id: 6

Next, lower item 5 from 60 to 35. Its parent is position 2, which now holds item 4 with key 10. Since 35 is not smaller than 10, nothing moves. The class still does a parent check, but it finds the rule already holds.

after id 5 to 35: [6:5, 3:20, 4:10, 0:50, 2:30, 5:35, 1:40]

Finally, poll the top. Item 6 leaves, and the last item, item 1 with key 40, moves to position 0. It sinks past the smaller child, item 4 with key 10, and then past item 5 with key 35. The result shows the sink:

pollId: 6 then [4:10, 3:20, 5:35, 0:50, 2:30, 1:40]

A heap with n items has about log2 n levels. For seven items, that is 3 levels. For a million items, it is about 20. So a decrease-key or a poll never does more than about 20 swaps for a million items. The search for the item costs one list read, pos[id], no matter how big the heap is.

The seven-item example runs with the class above saved in the same file. Add this code after the class, in place of the driver, and run it:

def seven_items():
    h = IndexedMinHeap(7)
    for item_id, k in enumerate([50, 40, 30, 20, 10, 60, 70]):
        h.insert(item_id, k)
    print("after inserts:", h.dump())
    h.decrease_key(6, 5)
    print("after id 6 to 5:", h.dump())
    print("root id:", h.peek_id())
    h.decrease_key(5, 35)
    print("after id 5 to 35:", h.dump())
    print("pollId:", h.poll_id(), "then", h.dump())
 
 
seven_items()

Using it for shortest paths

Dijkstra's method finds the shortest distance from one start point to every other point, when every road length is zero or more. The method is named after the computer scientist Edsger Dijkstra. It works like this. Start with distance 0 at the start and infinity everywhere else. Repeatedly take the unsettled point with the smallest distance, call it settled, and check each road out of it. If a road gives a shorter distance to a neighbour, lower that neighbour's distance.

Lowering a neighbour's distance is exactly decrease-key. The indexed heap holds every unsettled point, with its best-known distance as the key.

Take a small directed graph with five points, numbered 0 to 4. A directed road goes one way only. Each road is written as from, to, length. The roads are 0 to 1 with length 4, 0 to 2 with length 1, 2 to 1 with length 2, 1 to 3 with length 1, 2 to 3 with length 5, and 3 to 4 with length 3.

The graph is a dictionary, which maps each point to a list of (next point, length) pairs. math.inf is Python's infinity, so it compares as larger than any number. The code below goes under the class, in the same file:

def dijkstra_demo():
    n = 5
    adj = {
        0: [(1, 4), (2, 1)],
        1: [(3, 1)],
        2: [(1, 2), (3, 5)],
        3: [(4, 3)],
        4: [],
    }
    dist = [math.inf] * n
    dist[0] = 0
    pq = IndexedMinHeap(n)
    for v in range(n):
        pq.insert(v, dist[v])
    decrease_calls = 0
    while len(pq) > 0:
        u = pq.poll_id()
        print(f"settle {u} at {dist[u]}")
        for v, length in adj[u]:
            candidate = dist[u] + length
            if pq.contains(v) and candidate < dist[v]:
                dist[v] = candidate
                pq.decrease_key(v, candidate)
                decrease_calls += 1
                print(f"  lower {v} to {candidate}  heap {pq.dump()}")
    print("final distances:", dist)
    print("decrease-key calls:", decrease_calls)
 
 
dijkstra_demo()

The output is:

settle 0 at 0
  lower 1 to 4  heap [1:4, 4:inf, 2:inf, 3:inf]
  lower 2 to 1  heap [2:1, 4:inf, 1:4, 3:inf]
settle 2 at 1
  lower 1 to 3  heap [1:3, 4:inf, 3:inf]
  lower 3 to 6  heap [1:3, 4:inf, 3:6]
settle 1 at 3
  lower 3 to 4  heap [3:4, 4:inf]
settle 3 at 4
  lower 4 to 7  heap [4:7]
settle 4 at 7
final distances: [0, 3, 1, 4, 7]
decrease-key calls: 6

Read the lines from settle 2 at 1 down to see the point of the method. Point 1 first got distance 4, from the road 0 to 1. Later the search found a road from 2 that gives distance 3. The heap lowered point 1 from 4 to 3 and moved it to the top. The text inf in the heap output marks points that are not reached yet.

The final distances are 0 for point 0, 3 for point 1, 1 for point 2, 4 for point 3 and 7 for point 4. Those match what you get by checking the paths by hand: 0 to 2 to 1 costs 1 plus 2, which is 3.

Checking the class against a simple model

A class like this has many small steps that can go wrong. A good way to check it is to run random operations on it and, at the same time, keep a plain dictionary of the expected keys. After each operation, compare the two. The check used a remove method, which is the answer to exercise 3. In a run of 2,000 random rounds with 200 steps each, mixing insert, decrease-key, remove and poll, there were 400,000 checks and no failures. The run used a fixed random seed, so it can be repeated.

Common mistakes

Calling `decrease_key` with a larger number

decrease_key only lowers. A larger number raises an error with a clear message:

Traceback (most recent call last):
  File "/path/to/errs.py", line 7, in <module>
    h.decrease_key(0, 9)
    ~~~~~~~~~~~~~~^^^^^^
  File "/path/to/indexed_heap.py", line 49, in decrease_key
    raise ValueError(f"new priority {new_priority} is larger than {self.key[item_id]}")
ValueError: new priority 9 is larger than 4

If a value can go up, you need a different operation, which sifts down.

An id that is too big for the lists

With capacity 3, an id of 3 is out of range. The error comes from contains, because contains reads pos[id] first. Here is a small program that does it:

from indexed_heap import IndexedMinHeap
 
h = IndexedMinHeap(3)
h.insert(0, 10)
h.insert(3, 2)

The real output is:

Traceback (most recent call last):
  File "/path/to/trace_dup.py", line 5, in <module>
    h.insert(3, 2)
    ~~~~~~~~^^^^^^
  File "/path/to/indexed_heap.py", line 21, in insert
    if self.contains(item_id):
       ~~~~~~~~~~~~~^^^^^^^^^
  File "/path/to/indexed_heap.py", line 15, in contains
    return self.pos[item_id] != -1
           ~~~~~~~~^^^^^^^^^
IndexError: list index out of range

The paths and line numbers depend on where you save the files.

A negative id is not an error in Python

Python lists accept negative indexes. pos[-1] means the last item of the list, so an id of -1 quietly reads and writes the wrong slot. With capacity 3, insert(-1, 5) does not raise an error. Here is what it does:

from indexed_heap import IndexedMinHeap
 
h = IndexedMinHeap(3)
h.insert(-1, 5)
print("dump after insert(-1, 5):", h.dump(), "len", len(h))
print("pos:", h.pos, "key:", h.key)
print("contains(2):", h.contains(2))

Output:

dump after insert(-1, 5): [-1:5] len 1
pos: [-1, -1, 0] key: [0, 0, 5]
contains(2): True

Item 2 was never inserted, but contains(2) says it is in the heap, because the write for id -1 landed in slot 2. Check that ids are between 0 and n minus 1 before you use them.

Inserting an id that is already in the heap

The class checks for this:

Traceback (most recent call last):
  File "/path/to/errs2.py", line 6, in <module>
    h.insert(1, 2)
    ~~~~~~~~^^^^^^
  File "/path/to/indexed_heap.py", line 22, in insert
    raise ValueError(f"already in heap: {item_id}")
ValueError: already in heap: 1

If you want to change an item's key, call decrease_key, not insert.

Passing a position where an id belongs

Position 2 and id 2 are different things. The heap may still run, but the wrong item moves, and no error is raised if that id is in the heap. Always pass the id that you gave the item at insert time.

Forgetting to update `pos` in `_swap`

If pos is wrong, decrease_key looks in the wrong slot and moves some other item, or it does nothing useful. The fix is to keep _swap as the only place that moves items.

Forgetting to set `pos[id]` to -1 when an item leaves the heap

Then contains still says the item is inside, and decrease_key acts on an item that is no longer in the heap.

Using a plain `heapq` list with lazy deletion

A different approach is to skip the indexed heap, push a second copy of the item with the new key into a heapq list as a pair (key, id), and ignore old copies when they come out. This is called lazy deletion. It works, and it is simpler to write. The cost is that the list can grow to hold many stale copies. The indexed heap keeps exactly one copy of each item, so its size never grows past n.

Practice

Each exercise has its answer below it. Work the answer out on paper before you read it.

Exercise 1: insert and lower

Insert items A with key 7, B with key 4, C with key 9 and D with key 2, using the ids 0, 1, 2 and 3 in that order. Then call decrease_key(2, 1). Write the heap after the four inserts and after the decrease, in the id:key format.

Show the answer

After the four inserts the heap is [3:2, 1:4, 2:9, 0:7], which is D, B, C, A. After the decrease the heap is [2:1, 1:4, 3:2, 0:7]. C sits at position 2 under D, so it swaps once with D and moves to the top. That is the only move. Running the class with these inputs prints these two lines.

Exercise 2: why two different sifts

Explain in your own words why decrease_key only needs _sift_up, and why poll_id needs _sift_down.

Show the answer

Lowering a key makes one item smaller. The only rule it can break is the one with its parent, because the item may now be smaller than the parent above it. Its children were already bigger than the old key, so they are still bigger. So only _sift_up is needed. Poll moves the last item to the top. That item may be bigger than its children, so it must sink, which is _sift_down.

Exercise 3: remove any item

Write a remove(self, item_id) method that takes any item out of the heap in O(log n) time.

Show the answer

Put the item's slot at the end of the heap, shrink the size, and then fix the position of the item that moved into the gap. That item may need to move up or down.

def remove(self, item_id):
    if not self.contains(item_id):
        raise ValueError(f"not in heap: {item_id}")
    p = self.pos[item_id]
    self._swap(p, self.size - 1)
    self.size -= 1
    self.pos[item_id] = -1
    if p < self.size:
        moved = self.heap[p]
        self._sift_up(p)
        self._sift_down(self.pos[moved])

The moved variable matters. After _sift_up, the item may be at a new position, so the code asks self.pos[moved] where it went before it sinks. This method passed the 2,000-round random test described above.

Exercise 4: the same key

Call decrease_key(0, 4) on an item whose key is already 4. Does it raise an error? What happens to the heap?

Show the answer

It does not raise an error. The check is new_priority > self.key[item_id], and 4 is not larger than 4. The key is set to 4 again. _sift_up finds the parent is not bigger, so nothing moves. The heap is the same as before. For a heap of 0:4 and 1:9, the dump after the call is [0:4, 1:9].

Exercise 5: count the decrease-keys

On the five-point graph in the Dijkstra example, how many times is decrease_key called? Which point is settled last, and at what distance?

Show the answer

It is called 6 times. The output shows six lower lines. Point 4 is settled last, at distance 7.

Summary

An indexed heap keeps three lists: the heap order, the position of each id, and the key of each id.
Decrease-key finds the item in O(1) by its id, then moves it up in O(log n).
Every swap must update the position list for both items it moves.
Dijkstra's method uses decrease-key each time a shorter road to an unsettled point is found.
Use decrease_key to lower a key, and insert only for an id that is not in the heap yet.
Python lists accept negative indexes, so check that each id is between 0 and n minus 1.

Indexed heap with decrease-key in Python

Why this matters

Start with a shelf of labelled boxes

Three lists that agree

A worked trace by hand

The Python class

Reading the class

A larger example

Using it for shortest paths

Checking the class against a simple model

Common mistakes

Calling decrease_key with a larger number

An id that is too big for the lists

A negative id is not an error in Python

Inserting an id that is already in the heap

Passing a position where an id belongs

Forgetting to update pos in _swap

Forgetting to set pos[id] to -1 when an item leaves the heap

Using a plain heapq list with lazy deletion

Practice

Exercise 1: insert and lower

Exercise 2: why two different sifts

Exercise 3: remove any item

Exercise 4: the same key

Exercise 5: count the decrease-keys

Summary

Stuck? Ask the tutor

Calling `decrease_key` with a larger number

Forgetting to update `pos` in `_swap`

Forgetting to set `pos[id]` to -1 when an item leaves the heap

Using a plain `heapq` list with lazy deletion