Why this matters when you build with AI
An AI tool will often answer a search problem with a loop that looks right. It tries values from 0 to 10^9 one by one, or it uses a ternary search with while (lo < hi) and a midpoint that never moves. On a small sample the code passes, and on a large input it times out or never stops.
The bigger risk is the proof. A model that writes "the check is monotone, so binary search works" has not shown it. If the check is not monotone, binary search returns an answer with no error message. This lesson makes you write the proof for each check, and then test the check against a brute force on random inputs.
Every program here was compiled with g++ -O2 -std=c++17 -Wall -Wextra on the machine that wrote the lesson, using g++ 15.2.0. The random tests compare each result with a brute force, and the outputs below are copied from those runs. The memory figures are arithmetic, except the node counts, which the program printed. The Python sample was run with python3 -I. Your numbers may differ on other machines, but the node counts and test results should not.
The idea: keep every version of a tree
A segment tree over the values 0 to m - 1 stores counts. Each node says how many numbers in the array fall in its value range. To find the k-th smallest in a range [l, r) of the array, you need the counts for that range only. One way to get them is to build a tree for the prefix [0, r) and another for the prefix [0, l), then subtract. The difference of two counts is the count in the range.
That idea needs one tree per prefix, which is too much memory if you store them all. A persistent segment tree solves this with path copying. Each new element creates a new tree that shares almost all of its nodes with the previous tree. Only the nodes on one path from the root to a leaf are new.
Here is a small example with values 0 to 3, and the array [2, 0, 2].
Version 0 is the empty tree. Every node has count 0.
Version 1 adds the value 2. The path from the root to leaf 2 is copied, which makes 3 new nodes (root, the right child of the root, and leaf 2). Leaf 2 has count 1, and the root has count 1. Everything else still points to the empty tree's nodes.
Version 2 adds the value 0. The path to leaf 0 is copied, which makes 3 new nodes. The right child of the root is not touched, so version 2 shares it with version 1.
Version 3 adds the value 2 again. Another 3 nodes are copied.
So three inserts create 9 new nodes, and no version is changed after it is made. Old versions remain valid, which is what the word persistent means here.
Answering k-th smallest in a range
Compress the values first. Sort the distinct values of the array, and replace each value by its index in that sorted list. Now the tree has m leaves, where m is the number of distinct values. Call the compressed value c.
Let root[i] be the tree that holds the first i elements. Then root[0] is the empty tree, and root[i + 1] is root[i] with the element at position i inserted. The count of a node in root[r] minus the count of the same node in root[l] is the number of elements in a[l..r) whose compressed value lies in that node's range.
To find the k-th smallest (0-based k), start at the root and look at the left child. The number of elements in the left half of the range is count(left child in root[r]) - count(left child in root[l]). If k is smaller than that number, the answer is in the left half. Otherwise subtract that number from k and go right. Each step halves the value range, so a query takes O(log m) time.
Worked example with a = [3, 1, 4, 1, 5, 9, 2, 6]. The sorted distinct values are 1, 2, 3, 4, 5, 6, 9, so m = 7. The query [0, 8) with k = 3 covers the whole array. Its sorted values are 1, 1, 2, 3, 4, 5, 6, 9, and the element at index 3 is 3. The walk down the tree goes like this. At the root, the left half of the values (1, 2, 3 and 4) holds 5 elements of the whole array, which is more than 3, so the walk goes left. The left half is split again: the values 1 and 2 hold 3 elements, which is not more than 3, so the walk subtracts 3 from k, leaving 0, and goes right. The value 3 holds the one element needed, so the walk ends at compressed value 2, which is the value 3.
The C++ program
The code below builds the tree and answers the queries. It also prints the number of nodes and the answers for a few queries, including one query on the first 3 elements. That is a version query: root[3] is the tree of the prefix {3, 1, 4}, and it can be read at any time.
// Persistent segment tree: k-th smallest in a[l..r) (0-based k), with versions.
#include <bits/stdc++.h>
using namespace std;
struct PersistentTree {
int m; // number of distinct compressed values
vector<int> lc, rc, cnt; // node 0 is the empty tree
vector<int> root; // root[i] = tree of the first i elements
explicit PersistentTree(int m_) : m(m_) {
lc.push_back(0); rc.push_back(0); cnt.push_back(0);
}
// New tree = tree `prev` with one more element at compressed value pos.
int insert(int prev, int lo, int hi, int pos) {
int cur = (int)lc.size();
lc.push_back(lc[prev]); rc.push_back(rc[prev]); cnt.push_back(cnt[prev] + 1);
if (lo == hi) return cur;
int mid = (lo + hi) / 2;
if (pos <= mid) {
int child = insert(lc[prev], lo, mid, pos);
lc[cur] = child;
} else {
int child = insert(rc[prev], mid + 1, hi, pos);
rc[cur] = child;
}
return cur;
}
void build(const vector<int>& comp) { // comp[i] in [0, m)
root.assign(comp.size() + 1, 0);
for (size_t i = 0; i < comp.size(); i++)
root[i + 1] = insert(root[i], 0, m - 1, comp[i]);
}
// k-th smallest (0-based) among elements at indexes [l, r) of the array.
int kth(int l, int r, int k) const {
int u = root[l], v = root[r];
int lo = 0, hi = m - 1;
while (lo < hi) {
int mid = (lo + hi) / 2;
int left = cnt[lc[v]] - cnt[lc[u]];
if (k < left) { u = lc[u]; v = lc[v]; hi = mid; }
else { k -= left; u = rc[u]; v = rc[v]; lo = mid + 1; }
}
return lo; // compressed value
}
};
int main() {
vector<long long> a = {3, 1, 4, 1, 5, 9, 2, 6};
vector<long long> vals = a;
sort(vals.begin(), vals.end());
vals.erase(unique(vals.begin(), vals.end()), vals.end());
vector<int> comp;
for (long long x : a) comp.push_back((int)(lower_bound(vals.begin(), vals.end(), x) - vals.begin()));
PersistentTree t((int)vals.size());
t.build(comp);
printf("nodes used: %zu\n", t.lc.size() - 1);
printf("versions: %zu\n", t.root.size());
struct Q { int l, r, k; };
vector<Q> qs = {{0, 8, 0}, {0, 8, 3}, {2, 6, 1}, {4, 5, 0}, {1, 4, 2}};
for (auto q : qs)
printf("kth[%d,%d) k=%d -> %lld\n", q.l, q.r, q.k, vals[t.kth(q.l, q.r, q.k)]);
// Version query: the multiset of the first 3 elements {3,1,4}.
printf("version 3, smallest -> %lld\n", vals[t.kth(0, 3, 0)]);
printf("version 3, largest -> %lld\n", vals[t.kth(0, 3, 2)]);
return 0;
}Real output of that program, compiled with g++ -O2 -std=c++17:
nodes used: 31
versions: 9
kth[0,8) k=0 -> 1
kth[0,8) k=3 -> 3
kth[2,6) k=1 -> 4
kth[4,5) k=0 -> 5
kth[1,4) k=2 -> 4
version 3, smallest -> 1
version 3, largest -> 4Check these by hand. [2, 6) is 4, 1, 5, 9, whose sorted form is 1, 4, 5, 9, and index 1 is 4. [1, 4) is 1, 4, 1, whose sorted form is 1, 1, 4, and index 2 is 4.
Testing against a brute force
The random test below generates 2000 arrays with 1 to 30 values, and values that repeat often (small value range, negative numbers included). For each array it makes 60 random queries, sorts the slice, and compares the k-th element with the tree's answer. This gives 120000 checks. The output was:
all 120000 random k-th queries match brute forceThe test does not prove the code correct, but it covers duplicates, single-element ranges (r = l + 1), and every k. Most bugs in this structure show up on tiny arrays with many repeated values, so keep the generator small and ask for many runs.
Memory: estimate before you submit
Each insert creates one node per level of the tree, so about ceil(log2 m) + 1 nodes. The tree in the program stores three int values per node, which is 12 bytes. So the memory is about n * (ceil(log2 m) + 1) * 12 bytes, plus 4 bytes for each entry of root.
The program measured these node counts with random distinct values:
n=200000 nodes=3737856 (excluding null node 0)
depth for m=200000: 18
depth for m=1000000: 20For n = 200000, the node count is 3,737,856, and the arithmetic gives 44,854,272 bytes, about 45 MB, for the three arrays. The vectors grow by doubling, so the real capacity can be larger than the arithmetic. Reserve the capacity when you know n if the limit is tight.
For n = m = 10^6, the depth is 20, so each insert creates 21 nodes. That is 21,000,001 nodes, which is about 252 MB at 12 bytes each. Many judges allow 256 MB, so this solution may not fit. Check the memory limit of the problem before you write the code. If the estimate is over the limit, the options are a smaller node, a different structure such as a wavelet tree (look up its details before you use it), or an offline method from lesson 06. The mistake shows up only when the judge reports memory limit exceeded, so do the estimate first.
Binary search on the answer
Many problems ask for the smallest or largest value of a number that meets a condition. The number might be a capacity, a speed, a distance or a count. Binary search on the answer works when the condition is monotone: if the answer x works, then every larger x also works (or every smaller x works). Then you can split the search range in half and ask one question: "does mid work?"
The question is called the check. A check is correct only if you can prove it. The proof has two parts. First, the check is correct for one value. Second, the check is monotone. If you cannot write the second part, do not use binary search yet. Test the check on small inputs first.
Worked example, written as an original problem. There are n items with weights w_1, ..., w_n in a fixed order. You must cut them into at most m batches. Each batch is a run of consecutive items. A batch may hold at most K items, and its total weight must be at most C. Find the smallest C for which a cut exists.
The check for a given C is greedy. Start a batch, add items while the batch stays within both limits, and start a new batch when the next item does not fit. Count the batches. If the count is at most m, C works. If one item weighs more than C, C does not work.
Why greedy is correct for a fixed C: taking the longest possible first batch never hurts. Suppose some valid cut has a first batch that ends earlier than the greedy one. Move its boundary right to the greedy boundary. The first batch still fits the limits, because greedy stopped only where the next item did not fit. The rest of the cut still covers all the items, so no more batches are needed. By repetition, the greedy cut needs no more batches than any valid cut.
Why the check is monotone in C: this is the proof you must be able to write. Let C' >= C. Run greedy with C and with C'. Claim: after each batch, the greedy run with C' has covered at least as many items as the run with C. Induction: at the start both have covered 0 items. Suppose both have covered the same first j items, or the C' run has covered more. The batch starting at item j + 1 in the C run ends at some item e. Every item in that batch fits the limits with C, and C' >= C, so the same items fit with C'. The C' run starts its batch at item j + 1 or earlier and can extend at least to e. So the claim holds. When the C run has covered all items in b batches, the C' run has covered all items in at most b batches. So if C works, C' works. That is monotone.
With monotone checks, binary search on C finds the smallest value. The range is from the largest single weight (below it nothing fits) to the total weight. The total always works when a cut exists at all, which needs m * K >= n: then the count limit is the only limit that matters, and the greedy batches need at most ceil(n / K) batches. Each check costs O(n), so the total is O(n log(sum of weights)).
Here is the same algorithm in Python. It is short enough to read, and it runs quickly for this input:
Real output of that program:
answer: 13
C = 6 batches: item too big
C = 7 batches: 5
C = 8 batches: 4
C = 9 batches: 4Check by hand: with C = 13, the batches are [4, 2, 7], [1, 5, 3] and [6], which is 3 batches with at most 3 items each. With C = 12, the greedy batches are [4, 2], [7, 1], [5, 3] and [6], which is 4 batches, so 12 fails.
The C++ version below does the same work. Its test tries every set of cuts on 3000 random inputs with n up to 9.
// Binary search on the answer: split n items into at most m consecutive
// batches; each batch has total weight <= C and at most K items. Minimise C.
#include <bits/stdc++.h>
using namespace std;
// Greedy: fill each batch as far as the limits allow. Returns the batch count.
int batches_needed(const vector<long long>& w, long long C, int K) {
int count = 0;
long long sum = 0;
int items = 0;
for (long long x : w) {
if (x > C) return INT_MAX; // one item alone does not fit
if (count == 0 || sum + x > C || items + 1 > K) {
count++; // open a new batch
sum = 0;
items = 0;
}
sum += x;
items++;
}
return count;
}
// Smallest C in [lo, hi] with batches_needed(C) <= m. Monotone in C.
long long min_capacity(const vector<long long>& w, int m, int K) {
long long lo = *max_element(w.begin(), w.end());
long long hi = accumulate(w.begin(), w.end(), 0LL);
while (lo < hi) {
long long mid = lo + (hi - lo) / 2;
if (batches_needed(w, mid, K) <= m) hi = mid; // mid works, answer <= mid
else lo = mid + 1; // mid fails, answer > mid
}
return lo;
}
int main() {
vector<long long> w = {4, 2, 7, 1, 5, 3, 6};
int m = 3, K = 3;
printf("answer = %lld\n", min_capacity(w, m, K));
for (long long C = 6; C <= 9; C++) {
int b = batches_needed(w, C, K);
if (b == INT_MAX) printf("C=%lld item too big\n", C);
else printf("C=%lld batches=%d\n", C, b);
}
return 0;
}Real output of that program:
answer = 13
C=6 item too big
C=7 batches=5
C=8 batches=4
C=9 batches=4The brute force in the test tries all 2^(n-1) cut sets, keeps the ones with at most m batches and at most K items per batch, and takes the smallest largest-batch total. The test chooses m and K so that a cut exists (m * K >= n). Output on this machine:
all 3000 random cases match brute forceParallel binary search, in one paragraph
Sometimes many queries each have their own answer, and each answer is found by binary search, and each check is expensive. Parallel binary search runs all the searches together. In each round, every query has a midpoint. The algorithm sorts the queries by midpoint and sweeps through the input once, applying one change at a time and answering each query's check at the right moment. Then every query moves its range in half. After about log rounds, all answers are known, and the total work is one sweep per round. This lesson does not include code for it. Write a brute force check before you trust your version, because the sweep order is the part that is easy to get wrong.
Ternary search on unimodal functions
Ternary search finds the maximum of a function that rises and then falls. Pick two points m1 and m2 that split the range into thirds. If f(m1) < f(m2), the maximum is to the right of m1, so drop everything up to and including m1. Otherwise the maximum is at or left of m2, so drop everything after m2. Each step keeps about two thirds of the range, so the search takes O(log n) steps.
Use it only when the function is unimodal. A function is unimodal when there is one peak, the function never falls before the peak, and never rises after it. For integer arrays, the code below works when the only flat part is at the top:
// Ternary search on integers for the maximum of a unimodal array.
#include <bits/stdc++.h>
using namespace std;
// Returns an index of the maximum, assuming f rises, then falls.
// Correct when the only flat part is at the top (or there is none).
int ternary_max(const vector<int>& f) {
long long lo = 0, hi = (long long)f.size() - 1;
while (hi - lo > 2) {
long long m1 = lo + (hi - lo) / 3;
long long m2 = hi - (hi - lo) / 3;
if (f[m1] < f[m2]) lo = m1 + 1; // the peak is right of m1
else hi = m2; // the peak is at or left of m2
}
int best = (int)lo;
for (long long i = lo; i <= hi; i++) if (f[i] > f[best]) best = (int)i;
return best;
}
int main() {
// Strictly unimodal: the peak is at index 6 only.
vector<int> a = {1, 3, 4, 6, 8, 9, 12, 10, 7, 2};
int ia = ternary_max(a);
printf("strict: index %d value %d\n", ia, a[ia]);
// Flat part below the peak: values 2 at indexes 0..6, peak 7 at index 7.
vector<int> b = {2, 2, 2, 2, 2, 2, 2, 7, 3, 1};
int ib = ternary_max(b);
printf("plateau: index %d value %d (true max is 7 at index 7)\n", ib, b[ib]);
return 0;
}Real output:
strict: index 6 value 12
plateau: index 0 value 2 (true max is 7 at index 7)The second line is the failure. The function has a flat run below the peak, so when f(m1) == f(m2) the peak can be on either side, and the code drops the wrong part. A binary search on the slope (comparing f(mid) with f(mid + 1)) has the same weakness, because a flat step says nothing about the side of the peak. If flat parts can sit below the peak, look for another argument or another algorithm, and test it against a brute force.
The random tests make this concrete. They generate 20000 arrays with a single flat top (strict rise, flat run at the peak, strict fall), and 20000 arrays with flat runs allowed below the peak:
flat top only: 20000 correct, 0 wrong
flat below top allowed: 20000 cases, 657 wrongAn exhaustive search over unimodal arrays with lengths 3 to 6 and values 0 to 3 found a failure at length 4: the array 0 0 0 1. The code returns index 0 with value 0, but the maximum is 1 at index 3.
The tie cannot be resolved by a side rule. 0 2 2 2 2 2 2 7 also fails: the code returns index 1 with value 2, and the peak is at index 7. A flat run can hide the peak on either side, and every tie rule throws away one side, so some input always loses the peak.
Integer pitfalls
Three problems come up again and again.
The first is the midpoint that does not move. Binary search on the slope often uses this loop:
while (lo < hi) {
int mid = (lo + hi) / 2;
if (f[mid] < f[mid + 1]) lo = mid; // bug: should be mid + 1
else hi = mid;
}When hi = lo + 1, the midpoint is lo, so lo = mid changes nothing and the loop never ends. The sample below has a guard that stops the loop after 5 steps, so the output shows the stuck state. Without the guard, the program runs until someone stops it:
#include <bits/stdc++.h>
using namespace std;
// Pitfall: lo = mid never moves when hi = lo + 1, because mid == lo.
int main() {
vector<int> f = {1, 3, 2};
int lo = 0, hi = 2;
int steps = 0;
while (lo < hi) {
int mid = (lo + hi) / 2;
if (f[mid] < f[mid + 1]) lo = mid; // bug: should be mid + 1
else hi = mid;
if (++steps > 5) { printf("still looping, lo=%d hi=%d\n", lo, hi); return 0; }
}
printf("peak at %d\n", lo);
return 0;
}Real output:
still looping, lo=0 hi=1The second is the ternary step that makes no progress. When hi - lo is 1 or 2, the two points m1 and m2 can equal lo or hi, and the loop condition hi - lo > 2 stops the outer loop before that happens. Keep that condition, and scan the last few points by brute force, as the code above does.
The third is overflow in lo + hi. For int values near 10^9 the sum can pass the limit of int. Write lo + (hi - lo) / 2 instead.
Complexity
A persistent segment tree build costs O(n log m) time and memory. A k-th query costs O(log m) time. The search for the answer costs the check cost times log of the answer range. For the batch problem, that is O(n log(sum of weights)). Ternary search costs O(log n) evaluations of f for an array of length n.
Common mistakes
- Forgetting that
root[r]androot[l]must both be used. A common bug is to walk one tree only, which gives the k-th element of the prefix[0, r)and not of[l, r). The symptom on the sample is an answer that is too small forl > 0. - Using the uncompressed value range
[0, 10^9]as the tree range. Each insert then creates about 31 nodes instead of about 19, so 200000 inserts need about 74 MB instead of about 45 MB. Compress first. - Accepting a binary search without the proof. A check that is not monotone gives an answer that is sometimes too large and sometimes too small, and the sample may pass by chance. Test with random small inputs against a brute force, and include the cases where the answer is at the lower and upper bound.
- Using
lo = midin a binary search wheremidcan equallo. The loop does not end. Uselo = mid + 1orhi = mid - 1with the right midpoint rule. - Trusting ternary search on a function with flat parts below the top. The output is a valid index with a wrong value, so there is no crash to notice.
Hands-on exercise
Work in a folder of your own. Do these steps in order.
- Copy the persistent tree code above into
pst.cpp, compile it withg++ -O2 -std=c++17, and run it. Your output should match the output shown in this lesson. - Write a brute force that sorts each slice and reads index
k. Write a generator that makes 2000 arrays with values from-5to20and length 1 to 30, and 60 queries per array. Compare the two. Keep the test small so it runs in under a few seconds. - Change the tree so that it answers "how many elements in
a[l..r)are less than or equal tox", wherexis any integer. Test it against a brute force in the same way. - Use the warehouse search on your own input. Pick
n = 8weights between 1 and 9,m = 3, andK = 3. Compute the answer by hand, then by program, and check that they match. - Run the ternary code on
0 0 0 1and on0 1 1 1. Explain each result in one sentence.
Prompt for asking an AI tool to review your binary search, with what a good answer looks like:
Here is my binary search on the answer, with the check function below it.
Do three things. First, say whether the check is monotone in the parameter,
and give a proof, or give a counterexample. Second, list the smallest input
where the search returns a wrong answer, if one exists. Third, say which
lines of the search could loop forever. Do not rewrite the code yet.
[paste code]A good answer names a specific monotone argument or a specific counterexample. A weak answer says "looks correct" and gives no input. Check any counterexample it gives by running it against your brute force.
Practice
Each practice item has a hint. Work it out before you open the answer.
-
Memory estimate. You have
n = 1000000values, all distinct, in a persistent tree with no compression. Estimate the memory of the node arrays in bytes, using 12 bytes per node.Hint: find the depth first, then the nodes per insert.
Show the answer
The depth for
m = 10^6is 20 (the program printed it). Each insert creates one node per level, plus the leaf, so 21 nodes per insert. That is 21,000,000 nodes plus the null node. At 12 bytes each, the total is about 252,000,012 bytes, roughly 252 MB. The arithmetic is exact for these arrays, and the real capacity can be larger because vectors grow by doubling. -
Count query. Using the tree from this lesson, describe how to count the elements in
a[l..r)that are less than or equal to a valuex.Hint: the walk for k-th goes left or right based on counts. Here the decision is based on
x.Show the answer
Start at the root of
root[l]androot[r]with the value range[0, m - 1]. Lettbe the compressed index of the largest value that is less than or equal tox, or stop ifxis below all values. Walk down toward leaft. At each node, when the walk goes right, add the left child's count difference to the answer, then go right. When the walk goes left, go left without adding. At the end, add the count difference at the leaft. This takes O(log m) time. Test it against a brute force before you rely on it. -
Monotone proof. In the warehouse problem, prove that the check is monotone in
K(the maximum number of items per batch): ifKworks, thenK + 1works.Hint: compare the greedy batches for
Kand forK + 1, as the argument forCdoes.Show the answer
With
K + 1, the greedy rule allows a batch to hold one more item. Use the same induction as forC: after each batch, the greedy run withK + 1has covered at least as many items as the run withK. A batch in theKrun holds at mostKitems and stays withinC, so the same items fit theK + 1batch, and theK + 1batch can reach at least the same end. So theK + 1run needs at most as many batches. The check is monotone inK, so the smallest workingKcan be found by binary search, provided the check works for eachK. -
Ternary check. Find the integer
xin[0, 1000]that maximisesf(x) = -(x - 371)^2. Use the ternary search from this lesson, and state the result without running code.Hint: this function is strictly unimodal, so the code in this lesson works without changes.
Show the answer
The maximum is at
x = 371, wheref = 0. Every other value gives a negative number. The ternary search returns 371 because the function is strictly rising up to 371 and strictly falling after it. -
Find a failing array. Write a program that searches all arrays of length 5 with values 0 to 2, keeps only the ones that rise and then fall (non-strictly), and prints the first one where the ternary search returns a value below the maximum.
Hint: loop over all 3^5 arrays, and for each one check that it never falls before its peak and never rises after it. Print the returned value and the true maximum.
Show the answer
The first failing array in that loop order is
0 0 0 0 1. The code returns 0, but the maximum is 1 at index 4. Other orders of the loop may find a different first array, and any of them is a valid answer if the returned value is below the maximum.
Checklist
- Write the check first. Prove it is correct for one value, then prove it is monotone.
- Test the check on random small inputs against a brute force, including the smallest values and the values at the bounds.
- For a persistent tree, count nodes per insert from the depth, then multiply by the bytes per node before you write the code.
- Compress values before you build a value-indexed tree.
- Use
lo + (hi - lo) / 2for every midpoint. - Make sure every binary search loop moves
loorhion every step. - Use ternary search only when the function rises then falls with no flat part below the top.
- Keep the brute force and the generator in your repository next to the solution.
Summary
A persistent segment tree keeps one version of the tree for each prefix, and it creates only O(log m) new nodes per insert. Two versions subtracted give the counts for a range, which makes k-th queries O(log m). Binary search on the answer is correct when the check is monotone, and the proof is what makes it correct. Ternary search works on rising-then-falling functions, but a flat run below the top breaks it, and the midpoint and overflow traps can stop the program or give wrong answers.