Merge sort tree for range order queries in Python

Why it matters

A sum tree answers "what is the total of this range". A minimum tree answers "what is the smallest value in this range". Neither answers "how many values in this range are below 40". That question depends on the order of the values, and a single number per node cannot hold that order. This lesson shows a structure that keeps the order. Each node stores the sorted list of values in its segment, so a query can count the values below a bound by searching those lists.

The structure is called a merge sort tree because it is built by the merge step of merge sort. You may have seen merge sort in the sorting module. This lesson does not teach sorting. It uses one fact from it: two sorted lists can be joined into one sorted list in a single pass.

A count that a sum cannot give

Take the list [5, 2, 8, 6]. Suppose you want the number of values between index 0 and index 3 that are less than 6. The answer is 2, because 5 and 2 are less than 6. A sum tree stores 21 for the whole range, and that number does not help. A minimum tree stores 2, which says only that the smallest value is 2.

To answer the count, the search needs to know which values lie below a bound. A sorted list answers that quickly. If the list is [2, 5, 6, 8], the values below 6 are the ones before the first 6. A binary search (a search that halves a sorted list at each step) finds the position of 6 in the list, and the position is the count of values below it.

The merge sort tree puts a sorted list like this in every node. Each node covers a segment of the list, and its list holds exactly the values in that segment, in order.

The idea in a picture

Take the list [5, 2, 8, 6] and build the tree over the indexes 0 to 3. The root covers the whole list. Its two children cover [0..1] and [2..3]. Each of those splits again, down to the single items.

The root holds the sorted list of all four values, [2, 5, 6, 8]. The left child holds the sorted values of 5 and 2, which is [2, 5]. The right child holds the sorted values of 8 and 6, which is [6, 8]. The leaves hold one value each.

The tree has about n * log2(n) values in total, because each value appears once in every level. Each level splits the list into segments, so each level holds all n values in its own lists.

Building the lists by merging

The build follows the same shape as the build in the first segment tree lessons. A leaf stores its one value. An inner node builds its two children first, then combines their lists. The new step is the combine. Instead of adding two numbers, the node merges two sorted lists into one.

Merging is simple to do with two pointers. One pointer walks the left list, and one walks the right list. At each step, the smaller of the two values under the pointers goes into the output, and its pointer moves on. When one list runs out, the rest of the other list is copied over.

Here is the merge for [2, 5] and [6, 8]. The pointers start at 2 and 6. The 2 is smaller, so it goes out and the left pointer moves to 5. Then 5 is smaller than 6, so 5 goes out. The left list is empty. The rest of the right list, 6 and 8, is copied. The output is [2, 5, 6, 8].

Each merge reads every value in the two lists once, so the merge of a node costs time proportional to the size of its segment. Each level of the tree covers all n values once, so one level costs O(n), and the whole build costs O(n log n).

The code stores each node's list as a plain Python list. A Python list holds a reference to each number, and the number objects are shared between levels, because merging copies references and does not create new numbers. The tree still uses more memory than a C array would, and a later section compares the two.

def build(self, a, node, l, r):
    if l == r:
        self.tree[node] = [a[l]]
        return
    mid = (l + r) // 2
    self.build(a, 2 * node, l, mid)
    self.build(a, 2 * node + 1, mid + 1, r)
    self.tree[node] = merge(self.tree[2 * node], self.tree[2 * node + 1])

The merge function does the two-pointer walk described above. It builds the output list step by step, so it never needs a fixed size up front.

def merge(x, y):
    out = []
    i = j = 0
    while i < len(x) and j < len(y):
        if x[i] <= y[j]:
            out.append(x[i])
            i += 1
        else:
            out.append(y[j])
            j += 1
    out.extend(x[i:])
    out.extend(y[j:])
    return out

Answering a count query

The query uses the three-case search from the building and querying lesson. For a range [l..r], the search starts at the root. A node outside the range returns 0. A node fully inside the range does not recurse. A node partly inside asks both children.

The difference is what a fully inside node returns. Instead of its stored sum, it returns the number of values in its sorted list that are below the bound. That number comes from a binary search. The list is sorted, so the search finds the first position where the value is at least the bound, and that position is the count of values below the bound. Python's bisect module does this search for you.

The standard decomposition covers a query range with at most two fully inside nodes per level. Each one costs a binary search of about log2(n) steps. A query therefore costs about O(log^2 n) in total. For n equal to one million, log2(n) is about 20. Up to two fully covered nodes per level, over about 20 levels, each with a binary search of about 20 steps, gives at most about 800 steps. That is far fewer than the million steps of a loop over the range.

The public method takes an inclusive flag. With False, it counts values strictly below the bound. With True, it counts values at most the bound. The flag lets one method serve both kinds of query.

def count(self, node, nl, nr, l, r, x, inclusive):
    if r < nl or nr < l:
        return 0
    if l <= nl and nr <= r:
        if inclusive:
            return bisect_right(self.tree[node], x)
        return bisect_left(self.tree[node], x)
    mid = (nl + nr) // 2
    return (self.count(2 * node, nl, mid, l, r, x, inclusive)
            + self.count(2 * node + 1, mid + 1, nr, l, r, x, inclusive))

The two bisect functions are the same ones from the coordinate compression lesson. bisect_left gives the first index whose value is at least x, which is the count of values below x. bisect_right gives the first index whose value is greater than x, which is the count of values at most x.

Counting a range of values

A count of values between two bounds lo and hi, both included, is a difference. Count the values at most hi, then take away the values below lo. What remains is the values from lo to hi.

def count_in_range(self, l, r, lo, hi):
    if lo > hi:
        return 0
    return self.count(1, 0, self.n - 1, l, r, hi, True) - self.count(1, 0, self.n - 1, l, r, lo, False)

The guard lo > hi returns 0 for an empty range of values. Without it, the subtraction would give a negative number for some inputs.

Example 1: the lists in each node

The first example builds the tree for the four values [5, 2, 8, 6] and prints the list in each node. The class follows the pieces shown above, with a dump method added so the lists can be seen. The program also runs a few counting queries on a six-item list, and it checks 1000 random queries against a loop. Press Run to see the output.

The output is:

nodes for [5, 2, 8, 6]:
  node 1: [2, 5, 6, 8]
  node 2: [2, 5]
  node 3: [6, 8]
  node 4: [5]
  node 5: [2]
  node 6: [8]
  node 7: [6]
values in [1..4] that are less than 6: 2
values in [0..5] that are at most 6: 4
values in [0..5] between 3 and 7: 4
values in [2..3] between 5 and 100: 2
values in [0..5] between 9 and 4: 0
random queries checked: 1000, mismatches: 0
values stored across all nodes for n = 2000: 23952

The dump shows the tree from the picture. Node 1 is the root, with the whole sorted list. Node 2 covers indexes 0 and 1, so its list is [2, 5]. Node 3 covers indexes 2 and 3, so its list is [6, 8]. Nodes 4 to 7 are the four leaves. Node 4 holds the value 5 at index 0, node 5 holds the 2 at index 1, node 6 holds the 8 at index 2, and node 7 holds the 6 at index 3.

The count lines can be checked by hand. The range [1..4] of the second list is the values 2, 8, 6 and 3. Of these, 2 and 3 are less than 6, so the answer is 2. The range [0..5] between 3 and 7 holds 5, 6, 3 and 7, which is four values. The range [2..3] holds 8 and 6, and both are between 5 and 100, so the answer is 2.

The random check uses a loop as the reference answer. A query is correct only if the tree and the loop agree. Zero mismatches means every query agreed. The seed of the random generator is fixed at 42, so the same numbers appear on every run.

The last line counts the stored values. For n = 2000, the tree has 12 levels, and each level stores at most 2000 values. The total is 23952, which is a little less than 2000 * 12 because the lowest levels are not full in every branch. The total depends only on the size of the list, so this number is the same on every machine.

Example 2: a query that uses the bounds

The second example in the program shows one case that is easy to get wrong. The query count_in_range(0, 5, 3, 7) counts values from 3 to 7 inclusive. It is built from two calls. The first counts values at most 7, and the second counts values below 3.

On the list [5, 2, 8, 6, 3, 7], the values at most 7 are 5, 2, 6, 3 and 7, which is five values. The values below 3 are 2, which is one value. Five minus one is four, and the answer is four values: 5, 6, 3 and 7. The list in the root for this list is [2, 3, 5, 6, 7, 8]. The bisect_right for 7 is index 5, and the bisect_left for 3 is index 1. The subtraction gives 4.

The range [2..3] in the program holds the values 8 and 6. Both are between 5 and 100, so the count is 2. This is the same kind of query as before, only on a smaller range.

Example 3: checking against a loop

The third example tests the tree on 2000 random values and 1000 random queries. A loop over the range is the simplest correct method, so the program compares each tree answer with the loop answer. The test is a good habit for any tree structure. A tree can look correct on a small case and still break on a range that crosses a boundary between nodes.

The queries are random in three parts. The left end l is random, and the right end r is a random distance past it. The bounds lo and hi are random, and hi is lo plus a random amount up to 400000000. Because the bounds are chosen at random, some queries are empty and some cover most of the values. The mismatch count stays at zero when the tree is correct.

Common mistakes

The most common mistake is mixing up the inclusive and exclusive counts. A query for values at most x uses bisect_right. A query for values below x uses bisect_left. In a test run, the range of values between 3 and 7 in [5, 2, 8, 6, 3, 7] was computed as the values below 7 minus the values below 3. That counts only values from 3 up to 6, so the result was 3 instead of 4. The 7 itself is missing. The correct version uses count(l, r, hi, True) for the upper end.

The second mistake is building the lists inside the query. Each query should only read the lists. If a query merges lists or sorts them again, the time per query grows to O(n) or more, and the point of the structure is lost. The lists are built once, when the tree is created.

The third mistake is memory. Python lists are flexible, but each slot holds a reference of 8 bytes on a 64-bit machine, and a list has a little extra overhead. For n equal to one million, the tree holds about 21 million slots, which is about 168 megabytes for the references alone. A C-style array of 4-byte integers from the array module would use about half of that, and NumPy arrays would be smaller still. For a learning project, plain lists are the right choice. For large data, switch to a compact array type.

The fourth mistake is forgetting that the range is inclusive on both ends. The query [l..r] includes index r. A loop written as for i in range(l, r) leaves out the last item, and the tree and the loop will disagree in the random test.

Practice

1. For the list [5, 2, 8, 6], write the sorted list stored in node 3 and in node 7. Use the tree from Example 1.

Show the answer

Node 3 covers indexes 2 and 3, which hold 8 and 6. Its sorted list is [6, 8]. Node 7 is the leaf for index 3, which holds 6. Its list is [6].

2. Using the root list [2, 5, 6, 8] from the same list, count the values less than 6 by hand. Then say how a binary search gives the same answer in one step.

Show the answer

The values less than 6 are 2 and 5, so the count is 2. A binary search for 6 finds that the first index holding a value at least 6 is index 2. Index 2 is the count of values before it, so the answer is 2 in one step.

3. The list is [9, 1, 4, 4, 7]. Count the values in indexes 1 to 3 that are at most 4, first by reading the list and then by using the tree idea.

Show the answer

The values in indexes 1 to 3 are 1, 4 and 4. All three are at most 4, so the count is 3. The tree would split the range into nodes, search each node's sorted list with bisect_right for 4, and add the results. The total is the same 3.

4. A student stores each node's list unsorted and sorts it inside each query. Explain what this changes in the cost of a query, and say why the structure is then no better than a loop for small lists.

Show the answer

Sorting a list of size k costs about k * log2(k) steps. A query that sorts the lists of its nodes does work close to the size of the range again, which is the cost of a loop. The gain from a sorted list comes from keeping the list sorted before any query runs. For a small list, the extra bookkeeping makes the tree slower than a loop.

5. A merge sort tree is built over one million values. About how many references does it store in total, and roughly how much memory is that for the references alone on a 64-bit machine?

Show the answer

Each value appears once in each level, and a tree over one million values has about 21 levels. That gives about 21 million references. Each reference takes 8 bytes, so the references take about 168 megabytes. The number objects themselves are shared between levels, and the lists add a little overhead on top.

Summary

• A merge sort tree stores a sorted list of the values in each node, built by merging the two child lists.
• A count query visits the same nodes as a sum query, and in each fully covered node it does a binary search on that node's list.
• Use bisect_right for values at most a bound and bisect_left for values below a bound.
• A count between two bounds is the count at most the high bound minus the count below the low bound.
• Build the lists once, and test the tree against a loop on random queries.