Dynamic segment tree for huge ranges in Python

Why this matters

A segment tree answers questions such as "what is the sum of the values between position l and position r" very fast. The earlier lessons in this module used a tree stored in a list. That list has one slot for every node, and the number of slots grows with the size of the range. When the positions are small, for example 1 to 100000, that is fine. When the positions are up to one billion, the list would need billions of slots, and most computers cannot hold that much memory.

A dynamic segment tree solves this problem. It does not reserve space for every position. It creates a node only when a value is stored in the part of the range that node covers. A tree over one billion positions can then hold a few thousand values, and each insert adds at most 31 nodes. The rest of the range costs nothing.

This lesson builds that tree in Python. It assumes you have seen the basic segment tree idea, where each node covers a part of the range and has two children. If the words "node", "child" and "range query" are new to you, read them as follows. A node is one box of memory that stores a summary for a part of the range. A child is a node one level down that covers half of that part. A range query is a question about all positions from l to r, inclusive.

The problem with a plain list

Suppose you want to count how many times each customer ID has appeared, where IDs run from 1 to 1000000000. You want to answer "how many events came from IDs between 100 and 5000". The simplest tree is a segment tree stored in a list, with each leaf holding one ID.

A segment tree over n positions needs about 4 * n slots when it is stored in a list with the usual index rules. For n equal to one billion, that is four billion slots. Each slot is a reference of 8 bytes on a 64-bit machine, so the list alone would need about 32 gigabytes. A list of one billion references would still need 8 gigabytes. That is far more than an ordinary laptop can spare, so the plain tree is not practical here.

The idea: build only the path you need

Think of the range 1 to 8 as a small example. A normal segment tree would make 15 nodes for this range: one for the whole range, two for the halves, four for the quarters and eight for the single positions. A dynamic tree starts with an empty root and adds nodes as values arrive.

Here is the range split into halves at each step. Each box is a range.

[1, 8]              root, empty at first
  [1, 4]            left half
    [1, 2]
      [1, 1]
      [2, 2]
    [3, 4]
      [3, 3]
      [4, 4]
  [5, 8]            right half
    [5, 6]
      [5, 5]
      [6, 6]
    [7, 8]
      [7, 7]
      [8, 8]

When the value at position 5 is added, only the nodes on the path from the root to position 5 are created. Those are [1, 8], [5, 8], [5, 6] and [5, 5]. The other nodes do not exist. The path has one node per level, and a range of eight positions has four levels, so one value creates four nodes.

The same rule holds for a huge range. The root covers one billion positions. Each level halves the range, so the path from the root to any single position has about 30 halvings, which gives 31 nodes including the root. One insert creates at most 31 nodes. That is the key number for the memory budget, and we will use it again later in this lesson.

How a node stores its data

Each node needs three pieces of data.

total: the sum of all values stored in the part of the range this node covers.
left: the index of the child node that covers the left half, or 0 if no child has been created.
right: the index of the child node that covers the right half, or 0 if no child has been created.

We store nodes in three parallel lists, one for each field. A node is then just a number, its index. Using lists of plain numbers keeps the memory small. A class with attributes would also work, but each object carries extra overhead, and that overhead matters when a tree has millions of nodes.

Node 0 has a special meaning. It stands for "no node here". Its total is always 0, and the code never writes data into it. When a query reaches a child whose index is 0, it knows that nothing was ever stored in that part of the range, so the answer for that part is 0. This avoids creating nodes just to answer a question.

Python lists grow when you append to them, so the three lists start with one entry for node 0 and get a new entry for each node that is created. That means no fixed capacity has to be chosen in advance. The lists still hold one entry per node, so the memory is still set by how many nodes exist.

Adding a value, step by step

We will add value 5 to a tree over the range 1 to 8. Each step shows what happens at one node.

The root is node 0, which means the tree is empty. The code creates a new node, which gets index 1. This node is the root, and it covers [1, 8]. Its total becomes 1 because we added 1 at position 5.
The range [1, 8] splits at mid = 4. Position 5 is greater than 4, so it belongs to the right half. The left child is untouched. The right child is node 0, so the code creates node 2 for [5, 8], sets its total to 1 and stores its index in right of the root.
The range [5, 8] splits at mid = 6. Position 5 is at most 6, so it belongs to the left half. The code creates node 3 for [5, 6], adds 1 to its total and stores its index in left of node 2.
The range [5, 6] splits at mid = 5. Position 5 is at most 5, so it goes left. The code creates node 4 for [5, 5]. This is a leaf, a node that covers exactly one position, so the code adds 1 to its total and stops.

The tree now has four nodes, which matches the path length. If we add 7 next, the root and node 2 already exist, so only the nodes for [7, 8] and [7, 7] are new. Nodes that already exist are reused, so each new value creates only the part of the path that is missing.

Querying a range

A range query asks for the sum of all positions from l to r. The idea is the same as for a list-based segment tree. Start at the root. If the node's range is completely inside [l, r], return its total. If the node's range does not overlap [l, r] at all, return 0. Otherwise, ask both children and add their answers.

The rule for an empty child is the important part for this lesson. If the child index is 0, return 0 at once. The code does not create anything during a query. That keeps queries from making the tree bigger, and it keeps the answer correct, because an empty child holds only zeros.

Each query visits at most about four nodes per level, so it takes time proportional to the number of levels, about 31 steps for a range of one billion. The cost does not grow with the number of values stored.

The program

The following complete program builds the tree, runs a few queries and then checks the answers against a simple brute force method. A brute force method is the slow, obvious approach that is easy to trust. Here, the brute force method is a Python dictionary that maps each position to its total, and a range sum walks through every entry and adds the ones inside the range.

Python · editable

import random

LO = 1
HI = 1_000_000_000

class DynamicSegmentTree:
    def __init__(self):
        # Node 0 means "no node here". Its total is always 0.
        self.left = [0]
        self.right = [0]
        self.total = [0]
        self.root = 0

def nodes_used(self):
        return len(self.total) - 1

def add(self, pos, delta):
        self.root = self._add_at(self.root, LO, HI, pos, delta)

def _new_node(self):
        self.left.append(0)
        self.right.append(0)
        self.total.append(0)
        return len(self.total) - 1

def _add_at(self, node, lo, hi, pos, delta):
        if node == 0:
            node = self._new_node()
        self.total[node] += delta
        if lo == hi:
            return node
        mid = lo + (hi - lo) // 2
        if pos <= mid:
            self.left[node] = self._add_at(self.left[node], lo, mid, pos, delta)
        else:
            self.right[node] = self._add_at(self.right[node], mid + 1, hi, pos, delta)
        return node

def query(self, l, r):
        return self._query_at(self.root, LO, HI, l, r)

def _query_at(self, node, lo, hi, l, r):
        if node == 0 or r < lo or hi < l:
            return 0
        if l <= lo and hi <= r:
            return self.total[node]
        mid = lo + (hi - lo) // 2
        return self._query_at(self.left[node], lo, mid, l, r) + self._query_at(self.right[node], mid + 1, hi, l, r)

values = [5, 1_000_000_000, 999_999_999, 42, 42]
tree = DynamicSegmentTree()
for v in values:
    tree.add(v, 1)
print("nodes used after 5 inserts:", tree.nodes_used())
print("count in [1, 100]:", tree.query(1, 100))
print("count in [999999999, 1000000000]:", tree.query(999_999_999, 1_000_000_000))
print("count in [1, 1000000000]:", tree.query(1, 1_000_000_000))
print("count in [43, 998000000]:", tree.query(43, 998_000_000))

rnd = random.Random(7)
ops = 10000
big = DynamicSegmentTree()
truth = {}
ok = True
for i in range(ops):
    pos = 1 + rnd.randrange(HI)
    delta = rnd.randrange(10) - 3
    big.add(pos, delta)
    truth[pos] = truth.get(pos, 0) + delta
    if i % 50 == 0:
        a = 1 + rnd.randrange(HI)
        b = 1 + rnd.randrange(HI)
        l, r = min(a, b), max(a, b)
        expected = sum(s for p, s in truth.items() if l <= p <= r)
        if expected != big.query(l, r):
            ok = False
print("random check matches dict:", ok)
print("nodes used for", ops, "random inserts:", big.nodes_used())

The output is:

nodes used after 5 inserts: 68
count in [1, 100]: 3
count in [999999999, 1000000000]: 2
count in [1, 1000000000]: 5
count in [43, 998000000]: 0
random check matches dict: True
nodes used for 10000 random inserts: 177334

Read the first line first. The first insert created 31 nodes, one for each level of the path. The other four inserts reused the root and some upper nodes, and they created the remaining 37 nodes between them. The total of 68 is far below the one billion slots a plain list would need.

The line for [1, 100] returns 3 because the values 5, 42 and 42 lie in that range. The line for [43, 998000000] returns 0 because none of the five values is in that range. That answer comes from the empty child rule: whole parts of the range that have no node contribute 0 without being visited in detail.

The random check is the important one. It inserted 10000 random positions with small changes, some positive and some negative, and it compared 200 random range sums with the dictionary answer. All of them matched. The random generator is seeded with 7, so the same numbers appear on every run.

The last line shows the memory picture. The tree used 177334 nodes for 10000 inserts. Many inserts share the upper levels of the tree, so the number of nodes is less than the worst case of 31 per insert. In the worst case, every insert creates a new path of 31 nodes, which would be 310000 nodes for this run. Even that worst case is a small fraction of the one billion slots of a plain list.

A second example: setting balances

The first example added values. Many real problems need to set a value instead, for example "the balance of customer 17 is now 120". Setting is not the same as adding, because the old value must be removed first. A tree can only add changes to its totals, so the code must work out the change.

The fix is to keep the current balance of each customer in a dictionary, which maps a key to a value. When customer 17 is set to a new balance, the code reads the old balance, adds the difference to the tree and stores the new balance in the dictionary. The dictionary gives the current value of one customer quickly, and the tree gives the total over a range of customers.

The program below runs the same class with a set_balance function. Press Run to see the output.

Python · editable

LO = 1
HI = 1_000_000_000

class DynamicSegmentTree:
    def __init__(self):
        self.left = [0]
        self.right = [0]
        self.total = [0]
        self.root = 0

def nodes_used(self):
        return len(self.total) - 1

def add(self, pos, delta):
        self.root = self._add_at(self.root, LO, HI, pos, delta)

def _new_node(self):
        self.left.append(0)
        self.right.append(0)
        self.total.append(0)
        return len(self.total) - 1

def query(self, l, r):
        return self._query_at(self.root, LO, HI, l, r)

balance = {}
tree = DynamicSegmentTree()

def set_balance(customer_id, value):
    old = balance.get(customer_id, 0)
    tree.add(customer_id, value - old)
    balance[customer_id] = value

set_balance(17, 500)
set_balance(900_000_000, 250)
set_balance(17, 120)
set_balance(1_000_000_000, 75)
print("total in ids [1, 1000]:", tree.query(1, 1000))
print("total in ids [1, 1000000000]:", tree.query(1, 1_000_000_000))
print("total in ids [100000000, 999999999]:", tree.query(100_000_000, 999_999_999))
print("balance of id 17 is", balance[17])
print("nodes used:", tree.nodes_used())

The output is:

total in ids [1, 1000]: 120
total in ids [1, 1000000000]: 445
total in ids [100000000, 999999999]: 250
balance of id 17 is 120
nodes used: 87

Customer 17 was set to 500, then to 120. The first set added 500. The second set added 120 - 500 = -380, so the tree now stores 120 for that ID. The total over all IDs is 120 + 250 + 75 = 445. The total over the middle range holds only customer 900000000, with balance 250.

Common mistakes

Allocating the whole range. A line like [0] * 1_000_000_000 is the very problem this lesson solves. A list of one billion references takes about 8 gigabytes on a 64-bit machine, before any value is stored in it. Do not make a plain list for the whole range, even if it looks short in the code.

Forgetting the node 0 check in queries. If _query_at does not return 0 when node is 0, the code reads total[0], which is always 0, and left[0] and right[0], which are also 0. In this program that happens to give the right answer. The danger is in changing the code later. If someone writes to node 0 by mistake, all empty children start to share a total, and the answers become wrong without any error message. Keep node 0 untouched and check for it in every function that reads a child.

Updating before creating the node. In _add_at, the new node is created at the top of the function, before its total is changed. If the total update came first, self.total[node] += delta would add to node 0, the empty node that must stay at 0. The new node would then start with total 0, the value would be lost, and a query over that position would return 0 instead of the value you added.

Using a list of nodes that is too short. Node indexes must match the length of the three lists. The _new_node method appends to all three lists together, so they always stay the same length. If you create a node by setting an index that does not exist yet, Python raises an error:

totals = [0]
totals[1] += 5

Traceback (most recent call last):
  File "example.py", line 2, in <module>
    totals[1] += 5
    ~~~~~~^^^
IndexError: list index out of range

The fix is to append the new node to all three lists, as _new_node does, and only then use its index.

Using the wrong midpoint for a range that does not split evenly. The midpoint must be computed as lo + (hi - lo) // 2, and the children must be [lo, mid] and [mid + 1, hi]. Every position then belongs to exactly one child. Python integers never overflow, so (lo + hi) // 2 would also give the right answer here. Keep the form lo + (hi - lo) // 2 anyway, because it is the form that works in languages with fixed-size integers, and the code stays easy to move between languages.

Practice

Exercise 1. Build a dynamic tree over the range 1 to 8, the same small range as in the diagram. How many nodes are created by one insert at position 3? List the nodes by their range.

Show the answer

Four nodes are created. The path from the root to position 3 is [1, 8], then [1, 4] because 3 is at most mid = 4, then [3, 4] because 3 is greater than mid = 2 in [1, 4], then [3, 3] because 3 is at most mid = 3 in [3, 4]. Each level adds one node, and this range has four levels.

Exercise 2. The range is 1 to 1000000000. What is the largest number of nodes that one insert can create?

Show the answer

The answer is 31. Each step halves the range, and the length goes from one billion down to one position in about 30 halvings, because 2 to the power of 30 is 1073741824, which is at least one billion. The path has 31 nodes, one for each of the 31 levels from the root to a leaf. A single insert at position 5 in the first program creates exactly 31 nodes for this reason.

Exercise 3. Why does query return 0 when it reaches a child whose index is 0, instead of creating that child?

Show the answer

A child index of 0 means no node was ever created for that part of the range, and no value was ever added there, so every position in that part holds 0. Returning 0 gives the right answer. Creating nodes during a query would make the tree grow for no reason, and it would use memory that later inserts may need.

Exercise 4. After add(10, 3) and add(11, 4), what does query(1, 10) return? What does query(10, 11) return? Explain both answers.

Show the answer

query(1, 10) returns 3, because position 10 is inside the range and position 11 is not. query(10, 11) returns 7, because both positions are inside the range. The totals are the values added at each position, and each query covers only the positions from l to r, inclusive.

Exercise 5. A program will make 200000 inserts with the tree from this lesson. What is the worst-case number of nodes, and how much memory do the three lists use for that many nodes on a 64-bit machine? Count 8 bytes for each list slot and ignore the size of the number objects.

Show the answer

The worst case is 31 new nodes for each insert, so 200000 times 31 is 6200000 nodes, plus node 0. Each node has three list slots, one for left, one for right and one for total, so there are about 18600000 slots. At 8 bytes each, that is about 149 megabytes. Real runs are usually smaller, because inserts share the upper levels of the tree. The number objects are extra, and small numbers are shared by Python, so the real use depends on the values.

Summary

A dynamic segment tree creates a node only when a value is added to its part of the range, so it can cover positions up to one billion.
One insert creates at most 31 nodes for a range of one billion.
Node 0 means "no node here". Queries return 0 for it and never create nodes.
To set a value rather than add to it, keep the current values in a dictionary and add only the difference to the tree.
Keep the three lists the same length, by appending to all of them together when a node is created.

Dynamic segment tree for huge ranges in Python

Why this matters

The problem with a plain list

The idea: build only the path you need

How a node stores its data

Adding a value, step by step

Querying a range

The program

A second example: setting balances

Common mistakes

Practice

Summary

Stuck? Ask the tutor