Segment tree beats: cap a range while retaining its sum

Last updated: 5 Oct 20269 min read

tutorial

This segment tree supports range chmin: replace every capacity in a half-open interval by the smaller of itself and a limit. A node stores its sum, largest value, second distinct largest value, and count of the largest. If the limit is at least the maximum, nothing changes. If it lies strictly above the second maximum, only the current maxima change, so the sum can be corrected without visiting leaves. Otherwise the update descends. A deferred cap is pushed to children before partial work; a parent maximum is enough to constrain a child because the second-maximum condition held when the parent was capped. The program supports range sum and caps, not range add, assignment, or range chmax.

Operational case

A depot has daily capacity readings 47, 19, 83, 61, 29, and 103. Capping positions one through four at 47 reduces 83 and 61 and leaves the other readings alone, producing total 292. A later cap of positions zero through two at 23 yields first-three total 65 while the last three total 179. A cap equal to an existing maximum is a no-op. A limit below the second maximum cannot be applied to a whole mixed node by changing only the maximum count; the update must descend or its sum will be wrong.

Working Python program

class CapacityCapTree:
    def __init__(self, capacities):
        if not capacities:
            raise ValueError("at least one capacity is required")
        self.length = len(capacities)
        size = 4 * self.length
        self.maximum = [float("-inf")] * size
        self.second = [float("-inf")] * size
        self.max_count = [0] * size
        self.total = [0] * size
        self._build(1, 0, self.length, capacities)

    def _pull(self, node):
        left, right = node * 2, node * 2 + 1
        self.total[node] = self.total[left] + self.total[right]
        top = max(self.maximum[left], self.maximum[right])
        self.maximum[node] = top
        self.max_count[node] = sum(self.max_count[child] for child in (left, right)
                                   if self.maximum[child] == top)
        self.second[node] = max(self.second[child] if self.maximum[child] == top
                                else self.maximum[child] for child in (left, right))

    def _build(self, node, start, stop, capacities):
        if stop - start == 1:
            self.maximum[node] = self.total[node] = capacities[start]
            self.max_count[node] = 1
            return
        middle = (start + stop) // 2
        self._build(node * 2, start, middle, capacities)
        self._build(node * 2 + 1, middle, stop, capacities)
        self._pull(node)

    def _cap_node(self, node, limit):
        if self.maximum[node] <= limit:
            return
        assert self.second[node] < limit
        self.total[node] -= (self.maximum[node] - limit) * self.max_count[node]
        self.maximum[node] = limit

    def _push(self, node):
        for child in (node * 2, node * 2 + 1):
            if self.maximum[child] > self.maximum[node]:
                self._cap_node(child, self.maximum[node])

    def cap(self, left, right, limit):
        if not 0 <= left <= right <= self.length:
            raise IndexError("invalid half-open range")
        self._cap(1, 0, self.length, left, right, limit)

    def _cap(self, node, start, stop, left, right, limit):
        if right <= start or stop <= left or self.maximum[node] <= limit:
            return
        if left <= start and stop <= right and self.second[node] < limit:
            self._cap_node(node, limit)
            return
        self._push(node)
        middle = (start + stop) // 2
        self._cap(node * 2, start, middle, left, right, limit)
        self._cap(node * 2 + 1, middle, stop, left, right, limit)
        self._pull(node)

    def sum_range(self, left, right):
        if not 0 <= left <= right <= self.length:
            raise IndexError("invalid half-open range")
        return self._sum(1, 0, self.length, left, right)

    def _sum(self, node, start, stop, left, right):
        if right <= start or stop <= left:
            return 0
        if left <= start and stop <= right:
            return self.total[node]
        self._push(node)
        middle = (start + stop) // 2
        return (self._sum(node * 2, start, middle, left, right)
                + self._sum(node * 2 + 1, middle, stop, left, right))


if __name__ == "__main__":
    capacities = CapacityCapTree([47, 19, 83, 61, 29, 103])
    capacities.cap(1, 5, 47)
    print("after first cap:", capacities.sum_range(0, 6))
    capacities.cap(0, 3, 23)
    print("first three:", capacities.sum_range(0, 3))
    print("last three:", capacities.sum_range(3, 6))

Output

after first cap: 292
first three: 65
last three: 179

Time, space, and tradeoff

Build time and storage are O(N). A range-sum query visits O(log N) tree nodes after deferred caps are pushed. One cap can visit O(N) nodes in the worst case; repeated operations gain an amortized bound from the decreasing maximum structure, but that is not a per-operation latency promise. The Python implementation uses four arrays of size proportional to N and recursive calls of depth O(log N). It handles signed integer capacities, though a real capacity policy may reject negatives at its API boundary. A simple lazy segment tree is easier when updates are only additive.

Common Mistakes

Do not cap a node when the limit is at or below its second distinct maximum.
Do not update the sum without multiplying the drop by the maximum count.
Do not forget to push a parent cap before reading a partial child range.
Do not promise every individual cap is logarithmic.

Connected lessons

Compare its update and query contract with Palindromic trees: index distinct palindromes as text arrives, Wavelet matrices: count frequencies and find subarray quantiles, Two-stack window aggregation: keep FIFO order under a monoid, then complete the structure audit and decision quiz.

data structures

range-query-structures