Maximum subarray with node merging in Python

Why this matters

A subarray is a run of values that sit next to each other in a list. The maximum subarray problem asks for the run with the largest sum. Take daily profit changes of -2, 1, -3, 4, -1, 2, 1, -5, 4. The best run is the 4th to 7th values, which are 4, -1, 2, 1, and its sum is 6.

Kadane's algorithm solves this in one pass. It walks from left to right and keeps the best sum of a run that ends at the current value. That works well for one list that never changes. It does not work well when values change often, because each change means a new pass over every value.

A segment tree can answer the question for any range, and it can take an update in O(log n) time, which grows by one step each time n doubles. The difficulty is that a single node cannot store just one number. The best run inside a parent node may start in the left child, end in the right child, and cross the middle. To know that, a node has to remember more than one number about its children. This lesson uses four numbers per node. Two children combine into their parent with a fixed number of additions and comparisons, so each merge is O(1).

Four numbers for each segment

A segment tree is a binary tree where each node covers a range of positions. The root covers the whole list. Its two children cover the left half and the right half. A leaf covers one position. Each node stores a summary of its own range, and the summary is what makes merging possible.

For a range of values, the summary has four numbers. Every run used in these numbers must contain at least one value.

total is the sum of all values in the range.
prefix is the largest sum of a run that starts at the left end of the range.
suffix is the largest sum of a run that ends at the right end of the range.
best is the largest sum of any run inside the range.

A run must not be empty. This matters when every value is negative. The best run of -3, -1, -2 is -1, the single value -1. If empty runs were allowed, the answer would be 0, which is a different question.

Here is the summary for the range 2, -1, 3. Every run is listed so you can check each number.

The runs that start at the left end are 2, which sums to 2, 2, -1, which sums to 1, and 2, -1, 3, which sums to 4. The largest is 4, so prefix is 4.

The runs that end at the right end are 3, which sums to 3, -1, 3, which sums to 2, and 2, -1, 3, which sums to 4. The largest is 4, so suffix is 4.

Every run is 2, -1, 3, 2, -1 (sum 1), -1, 3 (sum 2) and 2, -1, 3 (sum 4). The largest is 4, so best is 4. The total is also 4.

A leaf holds one value v. Its runs are only v, so all four numbers equal v.

Merging two children

Call the left child L and the right child R. The parent P covers both ranges. Its four numbers come from the children with these rules.

The total is the simplest. P.total is L.total + R.total, because the sum of the whole range is the sum of its two halves.

The prefix of P is either a run that stays inside L, or a run that covers all of L and then continues into R. A run that stays inside L is one of the prefixes of L, so the best of those is L.prefix. A run that covers all of L and goes on into R has the sum L.total plus the best prefix of R, which is R.prefix. So P.prefix is the larger of L.prefix and L.total + R.prefix.

The suffix works the same way from the other side. P.suffix is the larger of R.suffix and R.total + L.suffix.

The best run of P is either inside L, inside R, or it crosses the middle. A crossing run is a suffix of L joined to a prefix of R. The best crossing run has the sum L.suffix + R.prefix. So P.best is the largest of L.best, R.best and L.suffix + R.prefix.

Example 1: merging `2, -1` with `3`

Start with the two halves of 2, -1, 3. The left half is 2, -1, and the right half is 3.

For the left half, L.total is 1, L.prefix is 2, L.suffix is 1 (the run 2, -1), and L.best is 2. The leaf for 3 has all four numbers equal to 3.

Now merge them.

total is 1 + 3, which is 4.
prefix is the larger of L.prefix = 2 and L.total + R.prefix = 1 + 3 = 4. The answer is 4.
suffix is the larger of R.suffix = 3 and R.total + L.suffix = 3 + 1 = 4. The answer is 4.
best is the largest of L.best = 2, R.best = 3 and L.suffix + R.prefix = 1 + 3 = 4. The answer is 4.

The merged node says the best run in 2, -1, 3 is 4. That matches the list of runs from the previous section. The crossing run is 2, -1, 3 itself, which starts in the left half and ends in the right half.

Each merge uses a fixed number of additions and comparisons. It does not matter whether the children cover 2 values or 2 million values. So a merge is O(1).

The costs of the three operations follow from this. Building the tree merges each internal node once, so it takes O(n) time. A point update changes one leaf and then recomputes the nodes on one path from the leaf to the root, which is about log n nodes, so it takes O(log n). A range query merges about 2 log n nodes that cover the range, so it takes O(log n) too.

Example 2: a full program with updates and queries

The program below builds the tree for -2, 1, -3, 4, -1, 2, 1, -5, 4, answers several range queries, changes one value, and asks again. It ends with a check against a slow brute-force method on 300 random lists.

The summary is a dataclass, a short class that only holds fixed fields. Each field can be read with a dot, such as info.best.

A brute-force method tries every possible run and keeps the best. It is slow, but it is simple enough to trust, so it serves as the reference answer.

The query method returns None when the query range does not overlap a node at all. This is the only case that returns None. Its caller then merges a None with a real summary, and merge returns the other side.

Python · editable

import random
from dataclasses import dataclass
from typing import Optional

@dataclass(frozen=True)
class Info:
    # The summary of one segment. Every field is about the segment only.
    total: int
    prefix: int
    suffix: int
    best: int

@staticmethod
    def leaf(value):
        return Info(value, value, value, value)

# Combine the summary of the left segment with the summary of the right segment.
    @staticmethod
    def merge(left: Optional["Info"], right: Optional["Info"]) -> Optional["Info"]:
        if left is None:
            return right
        if right is None:
            return left
        total = left.total + right.total
        prefix = max(left.prefix, left.total + right.prefix)
        suffix = max(right.suffix, right.total + left.suffix)
        best = max(left.best, right.best, left.suffix + right.prefix)
        return Info(total, prefix, suffix, best)

class SegmentTree:
    def __init__(self, a):
        self.n = len(a)
        self.tree = [None] * (4 * self.n)
        self._build(a, 1, 0, self.n - 1)

def _build(self, a, node, l, r):
        if l == r:
            self.tree[node] = Info.leaf(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] = Info.merge(self.tree[2 * node], self.tree[2 * node + 1])

def update(self, i, value):
        self._update(1, 0, self.n - 1, i, value)

def _update(self, node, l, r, i, value):
        if l == r:
            self.tree[node] = Info.leaf(value)
            return
        mid = (l + r) // 2
        if i <= mid:
            self._update(2 * node, l, mid, i, value)
        else:
            self._update(2 * node + 1, mid + 1, r, i, value)
        self.tree[node] = Info.merge(self.tree[2 * node], self.tree[2 * node + 1])

# Returns None when the query range does not overlap this node at all.
    def query(self, ql, qr):
        return self._query(1, 0, self.n - 1, ql, qr)

def _query(self, node, l, r, ql, qr):
        if qr < l or r < ql:
            return None
        if ql <= l and r <= qr:
            return self.tree[node]
        mid = (l + r) // 2
        return Info.merge(self._query(2 * node, l, mid, ql, qr), self._query(2 * node + 1, mid + 1, r, ql, qr))

# Reference answer: try every start and end pair. Slow, but obviously correct.
def brute_best(a, ql, qr):
    best = float("-inf")
    for i in range(ql, qr + 1):
        s = 0
        for j in range(i, qr + 1):
            s += a[j]
            best = max(best, s)
    return best

# Example 1: the whole list.
a = [-2, 1, -3, 4, -1, 2, 1, -5, 4]
st = SegmentTree(a)
whole = st.query(0, 8)
print(f"whole: best={whole.best} total={whole.total} prefix={whole.prefix} suffix={whole.suffix}")

# Example 2: updates and range queries.
print("query [0,3]: best=", st.query(0, 3).best)
print("query [4,8]: best=", st.query(4, 8).best)
st.update(7, 10)
print("after a[7]=10, query [4,8]: best=", st.query(4, 8).best)
print("after a[7]=10, query [0,8]: best=", st.query(0, 8).best)

# Check against the brute force on random lists, with a fixed seed.
rng = random.Random(42)
checks = 0
all_agree = True
for _ in range(300):
    n = 1 + rng.randrange(12)
    b = [rng.randrange(41) - 20 for _ in range(n)]
    t = SegmentTree(b)
    for _ in range(10):
        if rng.random() < 0.5:
            i = rng.randrange(n)
            v = rng.randrange(41) - 20
            b[i] = v
            t.update(i, v)
        else:
            ql = rng.randrange(n)
            qr = ql + rng.randrange(n - ql)
            checks += 1
            if t.query(ql, qr).best != brute_best(b, ql, qr):
                all_agree = False
print("random checks:", checks, "all agree:", all_agree)

The output is:

whole: best=6 total=1 prefix=2 suffix=5
query [0,3]: best= 4
query [4,8]: best= 4
after a[7]=10, query [4,8]: best= 17
after a[7]=10, query [0,8]: best= 20
random checks: 1490 all agree: True

Each line can be checked by hand. The whole list has a total of 1. The best prefix is -2, 1, -3, 4, -1, 2, 1, which sums to 2, so prefix=2. The best suffix is 4, -1, 2, 1, -5, 4, which sums to 5, so suffix=5. The best run is 4, -1, 2, 1, which sums to 6.

The query [0,3] covers -2, 1, -3, 4. The largest single run is 4, so the answer is 4. After a[7] becomes 10, the range [4,8] covers -1, 2, 1, 10, 4, and the best run is 2, 1, 10, 4, which sums to 17.

The random check counts 1490 queries. The exact count depends on the seed, so it is the same each time the program runs. The all agree: True result means every query matched the brute force.

Example 3: returning the positions of the best run

Sometimes the caller needs the run itself, or at least its start and end positions. The fix is to store more numbers in each node. prefix_end is the last position of the best prefix. suffix_start is the first position of the best suffix. best_start and best_end mark the best run.

The merge rules are the same as before. The positions just follow the winning choice. When the left prefix wins, prefix_end comes from the left child. When the right side wins, prefix_end comes from the right child, because the prefix now runs into the right child.

The following program uses a dataclass with eight fields and a merge that copies positions along with the sums. The tree is built the same way as before, except that a leaf is built with its position. The program checks that the reported positions lie inside the query, and that the values there add up to the reported best.

Python · editable

import random
from dataclasses import dataclass
from typing import Optional

@dataclass(frozen=True)
class Info:
    # Same idea as before, plus the positions where each best piece starts and ends.
    total: int
    prefix: int
    prefix_end: int
    suffix: int
    suffix_start: int
    best: int
    best_start: int
    best_end: int

@staticmethod
    def leaf(index, value):
        return Info(value, value, index, value, index, value, index, index)

@staticmethod
    def merge(left: Optional["Info"], right: Optional["Info"]) -> Optional["Info"]:
        if left is None:
            return right
        if right is None:
            return left

total = left.total + right.total

if left.prefix >= left.total + right.prefix:
            prefix, prefix_end = left.prefix, left.prefix_end
        else:
            prefix, prefix_end = left.total + right.prefix, right.prefix_end

if right.suffix >= right.total + left.suffix:
            suffix, suffix_start = right.suffix, right.suffix_start
        else:
            suffix, suffix_start = right.total + left.suffix, left.suffix_start

best, best_start, best_end = left.best, left.best_start, left.best_end
        if right.best > best:
            best, best_start, best_end = right.best, right.best_start, right.best_end
        crossing = left.suffix + right.prefix
        if crossing > best:
            best, best_start, best_end = crossing, left.suffix_start, right.prefix_end

return Info(total, prefix, prefix_end, suffix, suffix_start, best, best_start, best_end)

class SegmentTree:
    def __init__(self, a):
        self.n = len(a)
        self.tree = [None] * (4 * self.n)
        self._build(a, 1, 0, self.n - 1)

def _build(self, a, node, l, r):
        if l == r:
            self.tree[node] = Info.leaf(l, 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] = Info.merge(self.tree[2 * node], self.tree[2 * node + 1])

def query(self, ql, qr):
        return self._query(1, 0, self.n - 1, ql, qr)

a = [-2, 1, -3, 4, -1, 2, 1, -5, 4]
st = SegmentTree(a)
r = st.query(0, 8)
print("best sum", r.best, "from index", r.best_start, "to index", r.best_end,
      "(values", a[r.best_start:r.best_end + 1], ")")

rng = random.Random(7)
checks = 0
ok = True
for _ in range(300):
    n = 1 + rng.randrange(12)
    b = [rng.randrange(41) - 20 for _ in range(n)]
    t = SegmentTree(b)
    for _ in range(10):
        ql = rng.randrange(n)
        qr = ql + rng.randrange(n - ql)
        res = t.query(ql, qr)
        checks += 1
        # The reported piece must lie inside the query, and its sum must equal the best.
        inside = ql <= res.best_start <= res.best_end <= qr
        if not inside or sum(b[res.best_start:res.best_end + 1]) != res.best:
            ok = False
        brute = max(sum(b[i:j + 1]) for i in range(ql, qr + 1) for j in range(i, qr + 1))
        if brute != res.best:
            ok = False
print("random checks:", checks, "indices valid and sums match:", ok)

The output is:

best sum 6 from index 3 to index 6 (values [4, -1, 2, 1] )
random checks: 3000 indices valid and sums match: True

The positions point to values 4, -1, 2, 1, which are the positions 3 to 6. The program does not compare positions with a brute force, because two runs can have the same sum. It checks that the reported positions lie inside the query and that the values there add up to the reported best. It also checks that the best number matches the brute force.

Common mistakes

Allowing an empty run. If the leaf sets best to max(0, value), then an array of only negative values returns 0. The correct answer is the largest single value. For the array -3, -1, -2, every leaf would have best equal to 0, so the tree would answer 0. The correct answer is -1.

Leaving out the crossing run. If best is only the larger of left.best and right.best, the tree misses runs that span the middle. For 2, -1, 3, the answer should be 4, which is the crossing run 2, -1, 3. Without the crossing term, the tree would answer 3, the single value 3.

Asking for a range with the ends reversed. A query with ql greater than qr overlaps no node, so query returns None. Reading .best from that result fails. This is the error from a query query(2, 1) on a three-value tree:

AttributeError: 'NoneType' object has no attribute 'best'

The fix is to check that ql <= qr before the call, and to keep the bounds inside 0 and n - 1.

Swapping the children in merge. The rules are not symmetric. left.suffix + right.prefix gives a different answer from right.suffix + left.prefix. Swapping the two arguments produces a wrong answer and no error. Here the left half 2, -1 and the right half 3 are merged in the correct order, and then in the swapped order. This runs after the Info class from example 2 is defined:

left = Info.merge(Info.leaf(2), Info.leaf(-1))
print("correct merge", Info.merge(left, Info.leaf(3)))
print("swapped merge", Info.merge(Info.leaf(3), left))

The output is:

correct merge Info(total=4, prefix=4, suffix=4, best=4)
swapped merge Info(total=4, prefix=5, suffix=4, best=5)

The swapped answer says the best run has sum 5, which is not in the list of runs at all. The random check in example 2 compares only best, so a swap that happens to leave best unchanged can pass it. Test the prefix and suffix as well when you change the merge.

Using a number type that overflows. Python integers do not overflow, so this mistake does not happen here. In a language with fixed-size integers, a long sum can wrap around and give a wrong answer without an error. Keep that in mind when you port the code.

Practice

Find total, prefix, suffix and best for the list -4, 5, -2, by listing the runs.

Show the answer

total is -1. prefix is 1, from the run -4, 5. suffix is 3, from the run 5, -2. best is 5, from the single value 5. To check, build a tree for this list and print the fields of query(0, 2).

The left child has total 3, prefix 3, suffix 2 and best 4. The right child has total -2, prefix -1, suffix -2 and best -1. Find the four numbers of the parent.

Show the answer

total is 1, because 3 + (-2) is 1. prefix is 3, which is the larger of 3 and 3 + (-1) = 2. suffix is 0, which is the larger of -2 and -2 + 2 = 0. best is 4, which is the largest of 4, -1 and 2 + (-1) = 1.

Why must a leaf set best to its value, and not to max(0, value)?

Show the answer

The best run must contain at least one value. If every value in a list is negative, the best run is the least negative single value. Setting the leaf's best to 0 would claim that an empty run is allowed, so the answer for an all-negative list would be 0. The mistake section above shows this with the list -3, -1, -2.

Using the list 3, -2, 5, -1, what is the best sum for the range from index 1 to index 3?

Show the answer

The values in the range are -2, 5, -1. The runs are -2, 5, -1, 3 (for -2, 5), 4 (for 5, -1) and 2 (for -2, 5, -1). The best sum is 5, from the single value 5.

Use the version with positions. For the list -2, 1, -3, 4 and the range from index 0 to index 3, which positions does the best run cover?

Show the answer

The best sum is 4, and the run is the single value at index 3. Because 4 is larger than every other run in the range, the run covers index 3 alone, so best_start and best_end are both 3.

Summary

A node stores total, prefix, suffix and best for its range, and every run must contain at least one value.
Two children merge with a fixed number of additions and comparisons, so a merge takes O(1) time.
The crossing run is the best suffix of the left child joined to the best prefix of the right child.
Updates take O(log n) and queries take O(log n), which the tree gets from the same path walk as the earlier lessons.
The positions of the best run can be stored in the node and copied during the merge, at the cost of a larger node.
Check the tree against a brute-force method on random lists before you trust it.

Maximum subarray with node merging in Python

Why this matters

Four numbers for each segment

Merging two children

Example 1: merging 2, -1 with 3

Example 2: a full program with updates and queries

Example 3: returning the positions of the best run

Common mistakes

Practice

Summary

Stuck? Ask the tutor

Example 1: merging `2, -1` with `3`