Coordinate compression with a segment tree in Python

Why it matters

A segment tree needs a position for every item it stores. Positions are list indexes, so the tree needs a small range of whole numbers to work with. Real data often does not look like that. Prices, timestamps and ids can be large and spread out, and most of the range is empty. This lesson shows how to shrink such values to small positions while keeping their order, so the tree still answers questions about ranges of the original values.

The earlier lessons built trees over list positions, where position i is the item at index i. Here the position is a rank, which is the place a value takes in sorted order. The tree itself does not change. Only the way values are turned into positions is new.

The problem with large values

Suppose you want to know how many of the prices you have seen so far fall between 40 and 900. A tree indexed by price would need one leaf for every possible price. If prices go up to one billion, the tree needs about one billion leaves, and that is far more memory than a program should use for a handful of prices.

The fix is to notice that the tree only needs to know the order of the values. It does not need the values themselves to be small. If you have four distinct prices, you only need four leaves. The leaves can be numbered 0, 1, 2 and 3 in the order of the prices. A query for "prices from 40 to 900" becomes a query for "ranks from the first rank at or above 40 to the last rank at or below 900". That is a range over four small numbers.

This process has a name. Coordinate compression is the step of replacing large coordinates, here values, with their small ranks. The word coordinate means a position, so the step moves values to compact positions.

Step 1: sort and number the values

Take the values 900, 5, 42, 5, 1000000000, 42. Some of them repeat. The first job is to make a sorted list with each value appearing once.

In Python, a set removes repeats and sorted puts the rest in order:

values = [900, 5, 42, 5, 1000000000, 42]
sorted_distinct = sorted(set(values))

The result is [5, 42, 900, 1000000000]. The rank of a value is its index in this list, so 5 has rank 0, 42 has rank 1, 900 has rank 2 and 1000000000 has rank 3.

To find the rank of a value that is in the list, use list.index. For example, sorted_distinct.index(900) returns 2. The method needs the value to be present, and a later section shows what happens when it is not.

Now the four distinct values form a small tree with four leaves, even though the largest value is one billion. The tree stores counts or sums at ranks. Every query about the original values becomes a query about ranks.

Step 2: a tree over the ranks

The tree in this lesson has two operations.

The first is add(rank, delta). It adds delta to the value stored at one leaf, then fixes the sums on the path from that leaf up to the root. This is a point update, the same kind of update you met in the point updates lesson. The delta is 1 when you count a value and the value itself when you sum values.

The second is sum(l, r). It returns the total of the leaves from rank l to rank r, both included. This is the range query from the building and querying lesson.

The class below stores one total per leaf, and each inner node stores the sum of its two children. Python integers do not overflow, so the sums stay exact even when the values are large.

class RankTree:
    def __init__(self, size):
        self.size = size
        self.tree = [0] * (4 * size)
 
    def add(self, rank, delta):
        self._add(1, 0, self.size - 1, rank, delta)
 
    def _add(self, node, nl, nr, rank, delta):
        if nl == nr:
            self.tree[node] += delta
            return
        mid = (nl + nr) // 2
        if rank <= mid:
            self._add(2 * node, nl, mid, rank, delta)
        else:
            self._add(2 * node + 1, mid + 1, nr, rank, delta)
        self.tree[node] = self.tree[2 * node] + self.tree[2 * node + 1]
 
    def sum(self, l, r):
        if l > r:
            return 0
        return self._sum(1, 0, self.size - 1, l, r)
 
    def _sum(self, node, nl, nr, l, r):
        if r < nl or nr < l:
            return 0
        if l <= nl and nr <= r:
            return self.tree[node]
        mid = (nl + nr) // 2
        return self._sum(2 * node, nl, mid, l, r) + self._sum(2 * node + 1, mid + 1, nr, l, r)

The add method only goes down one path. At each level it picks the left or right child by comparing the rank with the middle. This is the same descent as the point update lesson, so the cost is O(log n) where n is the number of distinct values.

The sum method is the three-case query. A node outside the range returns 0, a node inside returns its stored total, and a node partly inside asks both children. The guard if l > r: return 0 covers a range that is empty after the bounds are converted, which happens when no value falls in the query.

The bounds problem

A query such as "values from 20 to 40" uses two numbers, 20 and 40, that may not be in the sorted list. The list is [5, 42, 900, 1000000000]. The value 20 is not there, and neither is 40. The tree needs ranks that match these bounds.

The bisect module, part of Python's standard library, has the two searches that are needed. bisect_left(sorted, x) returns the first index whose value is at least x. bisect_right(sorted, x) returns the first index whose value is greater than x. Both return an index from 0 to the length of the list. The values from x to y, both included, sit at the indexes from bisect_left(sorted, x) up to bisect_right(sorted, y) - 1.

Here is the list [5, 42, 900, 1000000000] with the bounds for three queries.

The third row shows why the guard in sum matters. Both bounds are past the end of the list, so the range is empty and the answer is 0.

Example 1: counting values in a range

The full program below is the first example in one file. It sorts and numbers the values, adds each one to a count tree and a sum tree, and then answers four queries. Press Run to see the output.

The output is:

sorted distinct: [5, 42, 900, 1000000000]
900 -> rank 2
5 -> rank 0
42 -> rank 1
5 -> rank 0
1000000000 -> rank 3
42 -> rank 1
[1, 100] count=4 sum=94
[42, 42] count=2 sum=84
[6, 899] count=2 sum=84
[2000000000, 2100000000] count=0 sum=0
keys: [10, 15, 20, 30, 500, 7000000]
item 0 is now 500; current values [500, 20, 30]
  items in [0, 100]: 2
  items in [15, 30]: 2
  items in [1000, 9000000]: 0
item 2 is now 15; current values [500, 20, 15]
  items in [0, 100]: 2
  items in [15, 30]: 2
  items in [1000, 9000000]: 0
item 1 is now 7000000; current values [500, 7000000, 15]
  items in [0, 100]: 1
  items in [15, 30]: 1
  items in [1000, 9000000]: 1

The first part of the output shows the ranks. The value 5 appears twice in the input, and both copies get rank 0. The second part answers the ranges. The range [1, 100] holds the values 5, 5, 42 and 42. There are four of them, and their sum is 94. The range [6, 899] holds only the two copies of 42, so the count is 2 and the sum is 84.

Example 2: a range with bounds that are not in the list

The second query in example 1 used [42, 42], where both bounds are in the list. The query [6, 899] is different. Neither bound, 6 or 899, is a stored value. The program still finds the right ranks.

Work through it by hand. The sorted list is [5, 42, 900, 1000000000]. The low bound for 6 is the first index whose value is at least 6. Index 0 holds 5, which is too small. Index 1 holds 42, which is large enough, so bisect_left returns 1. The high bound for 899 is the first index whose value is greater than 899. Index 2 holds 900, which is greater, so bisect_right returns 2. The last rank to use is 2 - 1, which is 1. The range of ranks is 1 to 1, and rank 1 is the value 42. The query counts two values and sums 84.

Now try the range [901, 999]. Its low bound is index 3, because 1000000000 is the first value at least 901. Its high bound is also index 3, because 1000000000 is greater than 999. The last rank is 3 - 1, which is 2. The range of ranks is 3 to 2. The left end is greater than the right end, so sum returns 0 at once. The range holds no values, and the guard handles it.

A plain rule covers all cases. Let frm be bisect_left of the low end, and let to be bisect_right of the high end, minus one. If frm is greater than to, the range is empty.

Example 3: values that change over time

The first two examples knew every value before they started. Real programs often do not. A value can change, and a new value can appear later. The safe approach is to read every value the program will ever use, including the new values in updates, and compress all of them at the start.

The example above starts with three items, 10, 20 and 30. The updates are: item 0 becomes 500, item 2 becomes 15, and item 1 becomes 7000000. Because the update values are known in advance, they are added to the list before the tree is built.

The program keeps the current value of each item in a list. For an update, it removes the old value from the tree by adding -1 at the old rank, then adds 1 at the new rank. Then it answers the three range queries. The live tree always holds exactly one count for each item, at the rank of its current value.

Notice the use of bisect_left(keys, current[item]) on the way out. The old value must be in keys, or the search returns the position where the value would sit, and the update goes to the wrong leaf. The compression step is what makes this safe. Every value that ever appears has a rank.

This approach is called offline compression. Offline means that the program reads all the data before it answers. If the values must be answered as they arrive, and a new value can show up that was not in the list, the program has to choose a different structure, such as a dynamic tree that builds nodes as it needs them. That structure is covered in the dynamic segment tree lesson.

Common mistakes

The most common mistake is using list.index on a value that is not in the list. Python raises an error, which is easy to notice. The mistake that is harder to notice is using bisect_left on a value that is missing and treating its result as the rank of that value. bisect_left returns the place where the value would sit, which may be the rank of a different stored value. Here the value 7 is missing from [5, 42, 900, 1000000000]:

keys = [5, 42, 900, 1000000000]
print(bisect_left(keys, 7))
print(bisect_right(keys, 7))

The output is:

1
1

The number 1 is the index of the value 42, which is the first value that is at least 7. It is not the rank of 7, because 7 has no rank. Use bisect_left and bisect_right for the bounds of a query. Use list.index only for a value that is known to be in the list. If you are not sure, check with value in sorted_distinct first.

Here is what list.index does with a missing value. The wording of the message can differ between Python versions:

keys = [5, 42, 900, 1000000000]
print(keys.index(7))

ValueError: list.index(x): x not in list

The second mistake is forgetting the - 1 after bisect_right. The high bound is the first index whose value is greater than the high end, so it is one past the last rank to include. In a test run, the query [42, 42] was written as counts.sum(bisect_left(sorted_distinct, 42), bisect_right(sorted_distinct, 42)) with no - 1. The result was count 3 and sum 984, because the value 900 at rank 2 was included. The correct answer is count 2 and sum 84.

The third mistake is sizing the tree from the raw values instead of the ranks. A tree over the raw values of one billion would need four billion slots. In a 64-bit Python list, each slot is a reference of 8 bytes, so the list would need 32000000000 bytes, which is about 32 gigabytes. The tree should use the number of distinct values, which is the length of sorted_distinct.

Practice

1. The values are 40, 10, 40, 1000. Write the sorted distinct list, and give the rank of each value in the order it appears.

Show the answer

The sorted distinct list is [10, 40, 1000]. The ranks in input order are 1, 0, 1, 2. The value 40 is at index 1, the value 10 is at index 0, and the value 1000 is at index 2.

2. Using the sorted list [10, 40, 1000], find bisect_left and bisect_right for x = 20, x = 40 and x = 2000. Then give the rank range for the values from 20 to 40.

Show the answer

For 20, bisect_left is 1 and bisect_right is 1. For 40, bisect_left is 1 and bisect_right is 2. For 2000, both are 3. The rank range for values from 20 to 40 is frm = bisect_left(20) = 1 to to = bisect_right(40) - 1 = 1. Only rank 1 is used, which is the value 40.

3. A student wants to count values from 100 up to the largest value in a sorted list keys, with no upper limit. Write the two numbers to pass to sum, and say what happens when no stored value is at least 100.

Show the answer

The left end is bisect_left(keys, 100), and the right end is len(keys) - 1. If no stored value is at least 100, bisect_left returns len(keys). That makes the left end larger than the right end, and the guard in sum returns 0. The answer is correct without a special case.

4. Explain in plain words why a tree sized to the number of distinct values is safe even when the input values are very large, and why the tree needs one leaf per distinct value and not one per possible value.

Show the answer

The tree only needs to keep the order of the values. Each distinct value gets one rank, and a query about a range of values becomes a query about a range of ranks. The number of distinct values is at most the number of items in the input, so the tree size depends on the input and not on the size of the numbers. A value of one billion takes the same space as a value of 5 once both have a rank.

Summary

• Coordinate compression replaces large values with their ranks in sorted order, so a tree can use small list positions.
• Use sorted(set(values)) to build the sorted list of distinct values.
• Use bisect_left for the first rank at or above a low bound and bisect_right - 1 for the last rank at or below a high bound.
• When values change over time, read all values first and compress them before building the tree.
• Never use a bisect result for a value that is missing from the list, because it is a position and not a rank.