Trees and Binary Search Trees Questions
Tree data structures as tested in coding interviews: binary trees, binary search trees, self-balancing trees (AVL insertion and deletion with rotations, a comparison with red-black trees, and alternatives such as treaps, splay trees and scapegoat trees), tries (prefix trees and radix trees), and N-ary, size-augmented and sum-augmented trees such as order-statistic trees. Covers traversal orders (in/pre/post-order, level-order, zigzag, iterative, generator-based and Morris variants), BST search, insertion and deletion invariants including duplicate-key policies, validation, successor and iterator design, range and kth-smallest queries, lowest common ancestor (including parent pointers, binary lifting and repeated queries on a static tree), tree reconstruction from traversals, serialization formats and parsing, converting between trees and sorted lists or linked lists, height, balance, diameter and path-sum recursion, subtree equality and hashing, persistent and snapshot-consistent trees, and how tree shape drives logarithmic versus degenerate performance. Excludes plain linked-list and stack/queue exercises with no tree in them, tree dynamic programming recurrences, general graph traversal, heap mechanics, and on-disk database index design, which are covered elsewhere.
Implement a weighted sampler with add(item, weight), update(index, new_weight) and sample(n) that returns items with probability proportional to weight. Use a binary sum tree (or Fenwick tree) so that update and sample are O(log n), explain how a sample descends the tree, and compare it with a prefix-sum array and with a heap.
Sample Answer
The idea in plain words
"Probability proportional to weight" means that if the weights are 3, 1, 4, 1, 5 (total 14), item e is picked with probability 5/14. A direct way to pick: draw a random whole number r from 0 to total - 1, lay the items out on a number line with item i owning a stretch of weight[i] consecutive positions, and return the item whose stretch contains r. The whole problem is finding that item fast while weights keep changing.
A binary sum tree does both parts in O(log n). Put the weights in the leaves of a binary tree (leaves are nodes with no children), and let every internal node store the sum of its two children. The root then holds the total. To find the owner of position r, start at the root and ask "is r inside the left child's stretch?": if r < left sum, go left; otherwise subtract the left sum from r and go right. One decision per level, so the walk takes as many steps as the tree has levels. To change a weight, add the difference to the leaf and to each of its ancestors, again one step per level. Heights and depths here count edges (empty tree -1, single node 0), so a tree with 8 leaves has height 3.
Weights here are non-negative integers, which makes the arithmetic exact; for decimal weights, scale them to integers first. Each update adds a difference to the ancestors, and with floating-point numbers the rounding can leave a tiny nonzero residue in a node whose true weight is 0, so a descent can then return a zero-weight item.
Brute force and the alternatives
- Linear scan: add weights until the running total passes r. Sampling is O(n), updates are O(1).
- Prefix-sum array (cumulative totals
[3, 4, 8, 9, 14]): sampling is O(log n) with a binary search, but changing one weight changes every cumulative value after it, so an update is O(n). - Heap: a heap keeps the largest (or smallest) item at the top and says nothing about sums of subtrees. It is the wrong tool for "pick in proportion to weight". (Heaps do appear in weighted sampling without replacement from a stream, where each item gets a random key built from its weight and a heap keeps the best k keys. That is a rarely asked extension and a different task from sampling from a changing table.)
- Sum tree or Fenwick tree: O(log n) for both operations.
The layout
The tree lives in a flat list. With size leaves (a power of two), leaf i is at index size + i, and the children of node j are 2j and 2j + 1. The five weights above, padded to 8 leaves, look like this (node numbers are the list indexes):
1:14
/ \
2:9 3:5
/ \ / \
4:4 5:5 6:5 7:0
/ \ / \ / \ / \
3 1 4 1 5 0 0 0
a b c d e
The root holds 14. The list is printed by the program below as [0, 14, 9, 5, 4, 5, 5, 0, 3, 1, 4, 1, 5, 0, 0, 0]; index 0 is unused.
Code
import random
class WeightedSampler:
"""Binary sum tree in an array. Leaf i is at size + i; each internal node holds the sum of
its two children, so tree[1] is the total weight. Weights are non-negative integers."""
def __init__(self, capacity=4):
self.size = 1
while self.size < capacity:
self.size *= 2
self.tree = [0] * (2 * self.size)
self.items = []
def add(self, item, weight):
if len(self.items) == self.size:
self._grow()
self.items.append(item)
self.update(len(self.items) - 1, weight)
def update(self, index, new_weight):
if not 0 <= index < len(self.items):
raise IndexError(index)
if new_weight < 0:
raise ValueError("weight must be >= 0")
node = self.size + index
delta = new_weight - self.tree[node]
while node >= 1: # add the change to every ancestor
self.tree[node] += delta
node //= 2
def _grow(self): # doubles the leaf level and rebuilds sums
leaves = self.tree[self.size:self.size + len(self.items)]
self.size *= 2
self.tree = [0] * (2 * self.size)
self.tree[self.size:self.size + len(leaves)] = leaves
for node in range(self.size - 1, 0, -1):
self.tree[node] = self.tree[2 * node] + self.tree[2 * node + 1]
def total(self):
return self.tree[1]
def find(self, r, log=None):
"""Index of the leaf that owns position r, for 0 <= r < total."""
node = 1
while node < self.size:
left = self.tree[2 * node]
if r < left:
node = 2 * node
where = "left"
else:
r -= left
node = 2 * node + 1
where = "right"
if log is not None:
log.append(f" at node {node // 2} (left sum {left}): go {where}, r is now {r}")
return node - self.size
def sample(self, n, rng):
if self.total() == 0:
raise ValueError("no weight to sample from")
return [self.items[self.find(rng.randrange(self.total()))] for _ in range(n)]
class Fenwick:
"""Fenwick (binary indexed) tree over the same weights: n + 1 slots instead of 2 * size."""
def __init__(self, weights):
self.n = len(weights)
self.t = [0] * (self.n + 1)
for i, w in enumerate(weights):
self.add(i, w)
def add(self, i, delta):
i += 1
while i <= self.n:
self.t[i] += delta
i += i & -i
def find(self, r):
"""Smallest index whose running total exceeds r."""
pos, step = 0, 1 << self.n.bit_length()
while step:
nxt = pos + step
if nxt <= self.n and self.t[nxt] <= r:
pos = nxt
r -= self.t[nxt]
step >>= 1
return pos
if __name__ == "__main__":
s = WeightedSampler()
for name, w in [("a", 3), ("b", 1), ("c", 4), ("d", 1), ("e", 5)]:
s.add(name, w)
print("capacity", s.size, "| tree array:", s.tree)
print("node 1 = total", s.total(), "| node 2 =", s.tree[2], "(leaves a..d) | node 3 =", s.tree[3], "(leaves e..h)")
log = []
idx = s.find(9, log)
print("find(9):")
print("\n".join(log))
print("-> index", idx, "item", s.items[idx])
# Exact check: over r = 0..total-1 every leaf must own exactly `weight` positions
def check(sampler):
owned = [0] * len(sampler.items)
for r in range(sampler.total()):
owned[sampler.find(r)] += 1
assert owned == sampler.tree[sampler.size:sampler.size + len(sampler.items)], owned
return owned
print("positions owned per item:", check(s))
s.update(1, 6) # b goes from 1 to 6
s.update(2, 0) # c goes from 4 to 0, so it must never be returned
s.add("f", 2)
s.add("g", 0)
s.add("h", 7)
s.add("i", 3) # the ninth item forces the tree to double
print("after updates, capacity", s.size, "| owned per item:", check(s))
fen = Fenwick(s.tree[s.size:s.size + len(s.items)])
assert all(fen.find(r) == s.find(r) for r in range(s.total()))
print("Fenwick find agrees with the sum tree on all", s.total(), "positions")
rng = random.Random(42)
draws = s.sample(20000, rng)
counts = {k: draws.count(k) for k in s.items}
print("20000 draws:", counts)
print("expected :", {k: round(20000 * w / s.total()) for k, w in zip(s.items, s.tree[s.size:])})
# Randomised check against a plain list of weights
rng = random.Random(9)
ref, t = [], WeightedSampler(1)
for step in range(400):
if not ref or rng.random() < 0.3:
ref.append(rng.randint(0, 5)); t.add(len(ref) - 1, ref[-1])
else:
i = rng.randrange(len(ref)); ref[i] = rng.randint(0, 5); t.update(i, ref[i])
if sum(ref):
assert check(t) == ref
print("random add/update sequence checked against a plain list: 400 steps, final size", len(ref))
# The two guards: nothing to sample, and a negative weight
for label, action in [("empty sampler, sample(1)", lambda: WeightedSampler().sample(1, rng)),
("update(0, -1) on the sampler above", lambda: s.update(0, -1))]:
try:
action()
except ValueError as err:
print(f"{label}: ValueError: {err}")
Output:
capacity 8 | tree array: [0, 14, 9, 5, 4, 5, 5, 0, 3, 1, 4, 1, 5, 0, 0, 0]
node 1 = total 14 | node 2 = 9 (leaves a..d) | node 3 = 5 (leaves e..h)
find(9):
at node 1 (left sum 9): go right, r is now 0
at node 3 (left sum 5): go left, r is now 0
at node 6 (left sum 5): go left, r is now 0
-> index 4 item e
positions owned per item: [3, 1, 4, 1, 5]
after updates, capacity 16 | owned per item: [3, 6, 0, 1, 5, 2, 0, 7, 3]
Fenwick find agrees with the sum tree on all 27 positions
20000 draws: {'a': 2144, 'b': 4418, 'c': 0, 'd': 760, 'e': 3728, 'f': 1477, 'g': 0, 'h': 5284, 'i': 2189}
expected : {'a': 2222, 'b': 4444, 'c': 0, 'd': 741, 'e': 3704, 'f': 1481, 'g': 0, 'h': 5185, 'i': 2222}
random add/update sequence checked against a plain list: 400 steps, final size 139
empty sampler, sample(1): ValueError: no weight to sample from
update(0, -1) on the sampler above: ValueError: weight must be >= 0
How a sample descends: the trace
find(9) on the tree above (the cumulative totals are 3, 4, 8, 9, 14, so position 9 is the first position owned by e):
- At node 1 the left child (node 2) holds 9. r = 9 is not below 9, so go right and set r = 9 - 9 = 0.
- At node 3 the left child (node 6) holds 5. r = 0 < 5, so go left.
- At node 6 the left child (leaf 12) holds 5. r = 0 < 5, so go left. Node 12 is a leaf: index 12 - 8 = 4, which is
e.
That is exactly the three lines the program prints, and the result e matches the layout.
Why it is correct, and checked exhaustively. At each node the left child owns the first left sum positions of that node's stretch and the right child owns the rest, so the decision never loses a position. The check function in the program tries every r from 0 to total - 1 and confirms that item i is returned for exactly weight[i] of them. That exact count (not a random test) is shown after the updates too, where a weight of 0 is owned by zero positions and so can never be sampled. The 20,000-draw run is a fixed-seed illustration: its counts land near the expected counts but are not exact, which is normal for random draws.
Add, update, and sample(n)
- update(index, new_weight): compute the difference, then add it to the leaf and to every ancestor up to the root: one step per level.
- add(item, weight): put the item in the next free leaf and call update. When the leaves are all used, the tree doubles its leaf count and rebuilds every internal sum in O(size). Amortized means the cost averaged over a whole run of operations. A rebuild at capacity 8 costs about 16 steps, but it can only happen again after 8 more adds have filled the new space, so the rebuild work per add stays constant on average even though one unlucky add costs O(n). In the program, the ninth add doubles the capacity from 8 to 16.
- sample(n): n independent draws of r, each followed by a descent: O(n log N) for N items. That is sampling with replacement (the same item can appear twice). Sampling without replacement can be done by setting a drawn item's weight to 0 with
update, and restoring it afterwards, O(log N) each.
Fenwick tree comparison
A Fenwick tree (binary indexed tree) stores partial sums in one list of N + 1 slots, using the lowest set bit of the index to decide which range each slot covers (worked in the walk-through below). Updates add the difference to O(log N) slots; the sample step finds the smallest index whose running total exceeds r by testing powers of two from large to small. The program builds one from the same weights and shows that its find returns the same index as the sum tree for all 27 positions. Choose the Fenwick tree when memory matters or you only need prefix sums plus sampling: it uses N + 1 slots instead of 2 * size (roughly 2N to 4N). Choose the explicit sum tree when you want to read a node's subtree total directly, such as to descend to a subtree, to split sampling by region, or to teach and debug.
Fenwick walk-through on the five weights
Three ideas make the Fenwick tree work. The lowest set bit of a number is its rightmost 1 in binary, and i & -i isolates it (for 12 = 1100 it is 4). Slot i (counting from 1) stores the sum of the i & -i weights that end at index i, so slot 4 covers four weights, slot 6 covers two, and every odd slot covers just itself. Adding i & -i to i jumps to the next slot whose range contains i, which is how an update reaches every slot that must change. And 1 << n.bit_length() is the smallest power of two above n (8 for n = 5), the largest jump the descent can try.
weights = [3, 1, 4, 1, 5]
n = len(weights)
t = [0] * (n + 1) # slot i (1-based) holds the sum of the last (i & -i) weights ending at i
def add(i, delta):
i += 1 # convert to 1-based
while i <= n:
t[i] += delta
print(f" slot {i} ({i:03b}) += {delta} -> {t[i]}")
i += i & -i # i & -i is the lowest set bit of i; adding it jumps to the next slot that covers i
for idx, w in enumerate(weights):
print(f"add weight {w} at index {idx}:")
add(idx, w)
print("slots 1..5:", t[1:])
for i in range(1, n + 1):
low = i & -i
print(f"slot {i} ({i:03b}): lowest set bit {low}, covers weights at indexes {i - low}..{i - 1} = {sum(weights[i - low:i])}")
def find(r):
pos, step = 0, 1 << n.bit_length()
print(f"find({r}): start pos=0, step={step}")
while step:
nxt = pos + step
if nxt <= n and t[nxt] <= r:
print(f" step {step}: slot {nxt} holds {t[nxt]} <= r={r}, take it, r becomes {r - t[nxt]}")
pos, r = nxt, r - t[nxt]
else:
print(f" step {step}: slot {nxt} " + ("is past the end" if nxt > n else f"holds {t[nxt]} > r={r}, skip"))
step >>= 1
return pos
print("owner of position 9 is index", find(9))
# growing a sum tree from 2 leaves to 4 when a third weight arrives
size, tree = 2, [0, 4, 3, 1] # leaves are tree[2:4] = [3, 1]; tree[1] = 4 is the total
print("\nbefore grow:", tree)
leaves = tree[size:]
size *= 2
tree = [0] * (2 * size)
tree[size:size + len(leaves)] = leaves # copy the leaves into the new, wider leaf level
for node in range(size - 1, 0, -1): # recompute every internal sum bottom-up
tree[node] = tree[2 * node] + tree[2 * node + 1]
print("after grow: ", tree)
tree[size + 2] = 4 # now the third weight fits; add it to the leaf (node 6) and its ancestors (nodes 3 and 1)
tree[3] += 4
tree[1] += 4
print("after adding weight 4 at index 2:", tree)
Output:
add weight 3 at index 0:
slot 1 (001) += 3 -> 3
slot 2 (010) += 3 -> 3
slot 4 (100) += 3 -> 3
add weight 1 at index 1:
slot 2 (010) += 1 -> 4
slot 4 (100) += 1 -> 4
add weight 4 at index 2:
slot 3 (011) += 4 -> 4
slot 4 (100) += 4 -> 8
add weight 1 at index 3:
slot 4 (100) += 1 -> 9
add weight 5 at index 4:
slot 5 (101) += 5 -> 5
slots 1..5: [3, 4, 4, 9, 5]
slot 1 (001): lowest set bit 1, covers weights at indexes 0..0 = 3
slot 2 (010): lowest set bit 2, covers weights at indexes 0..1 = 4
slot 3 (011): lowest set bit 1, covers weights at indexes 2..2 = 4
slot 4 (100): lowest set bit 4, covers weights at indexes 0..3 = 9
slot 5 (101): lowest set bit 1, covers weights at indexes 4..4 = 5
find(9): start pos=0, step=8
step 8: slot 8 is past the end
step 4: slot 4 holds 9 <= r=9, take it, r becomes 0
step 2: slot 6 is past the end
step 1: slot 5 holds 5 > r=0, skip
owner of position 9 is index 4
before grow: [0, 4, 3, 1]
after grow: [0, 4, 4, 0, 3, 1, 0, 0]
after adding weight 4 at index 2: [0, 8, 4, 4, 3, 1, 4, 0]
Read the trace like this. Adding weight 3 at index 0 touches slots 1, 2 and 4: the three slots whose ranges include index 0. The table shows slot 4 holds 9 = 3 + 1 + 4 + 1, the first four weights. In find(9) the descent tries the largest jump first: slot 8 does not exist, slot 4 holds 9 which is not above r = 9, so the whole block belongs before the answer and r drops to 0. The next jump lands on slot 6, past the end, and slot 5 holds 5, which is above 0, so the walk stops with pos = 4: index 4, item e, the same owner the sum tree found. The last three lines show the sum tree's growth step on a 2-leaf tree: the leaves are copied into a wider leaf level and every internal sum is recomputed bottom-up.
A typical interview needs one of the two trees (the sum tree or the Fenwick tree), the update rule and the descent rule. The alias method and sampling with random keys and a heap are deeper material that a coding round rarely asks for.
Complexity
- update, find: O(log size) each, where the tree has about log2(size) levels (3 for 8 leaves, 4 for 16).
- add: O(log size) plus an O(size) rebuild each time the capacity doubles: amortized O(log size).
- sample(n): O(n log size).
- Space: a list of 2 * size entries, where size is at most about 2N for N items, so about 4N entries at most.
Edge cases and common mistakes
- Total weight 0: there is nothing to sample; the code raises
ValueErrorinstead of returning an item (the last lines of the first program's output show it). - Zero-weight items must never be returned: use
r < left(strict) when choosing the left child, which the exhaustive check confirms. - Updating a leaf without updating its ancestors leaves every sum above it stale, so the descent is steered by old totals.
- Using floating-point weights and
random.random() * total: rounding can send the descent into a zero-weight leaf. Scale to integers or add a guard. - Off-by-one on the position:
randrange(total)yields 0 to total - 1; usingrandint(0, total)can pick a position past the end. - Negative weights make the stretches meaningless; reject them as
updatedoes (also shown in the last output line of the first program).
Follow-ups and how the answer changes
- Weights never change and the number of draws is huge: the alias method (a precomputed table of two-way coin flips; rarely asked in interviews) gives O(1) per draw after an O(n) setup, but any update means rebuilding.
- Sampling without replacement while weights change: the zero-and-restore trick above, in O(log N) per draw.
- Different sample sizes per call and a stream of items: keep the same tree and only grow the leaf level when needed.
Running the code
# Python 3.12, standard library only
python weighted_sampler.py # the sampler (first program)
python fenwick_walk.py # the Fenwick walk-through (second program)
That is every published Trees and Binary Search Trees question for Machine Learning Engineer so far. Browse the other topics in this category, or practice this one interactively.