Lab 13 — Treap & Ordered-Set Operations

Goal

Implement a treap — the one balanced BST you can realistically write in an interview. Build it from exactly two primitives, split and merge, then compose insert, delete, k-th smallest, rank, and count-less-than on top. Understand the implicit treap (position-keyed) for array cut/paste/reverse. Apply it to LC 315. After this lab: split + merge from memory in <10 minutes, the full ordered set in <25.

Background Concepts

A treap stores a key (BST-ordered) and a priority (max-heap-ordered) per node. For distinct priorities the shape is unique: the max-priority node is the root, recurse on both sides. That shape is exactly what you’d get by inserting keys into a plain BST in decreasing-priority order — so random priorities simulate random insertion order, giving expected depth O(log n) regardless of key order.

Everything reduces to two primitives:

  • split(t, key) → (left, right) where left holds keys < key, right holds keys ≥ key.
  • merge(a, b) → one treap, precondition: every key in a ≤ every key in b.

Insert = split + merge in a new node. Erase = unlink + merge children. Order statistics (k-th, rank) come from an augmented size field maintained by a pull() after every child mutation.

Implicit treap: drop explicit keys; a node’s “key” is its position, computed implicitly as size(left subtree) along the search path. Split by size instead of key. Now you can cut out any array range in O(log n), paste it elsewhere, or reverse it with a lazy flip flag — a rope.

Why this matters: C++ has std::set plus pbds order-statistics trees; Java has TreeMap (no rank); Python has no ordered set at all — sortedcontainers is third-party and banned on some judges. In Python, treap knowledge is the difference between solving rank-query problems and not.

Interview Context

“Implement a balanced BST” is rare, but ordered-set-with-rank problems are constant: LC 315, LC 327, LC 493, sliding-window medians, online percentile queries. AVL and red-black trees are unwritable in 30 minutes (rotation case analysis); a treap is ~70 lines with two primitives, and interviewers accept it. Saying “I’ll use a treap because split/merge composes cleanly” is a strong L5 signal; deriving the E[depth] bound on request seals it.

Problem Statement (LC 315)

Count of Smaller Numbers After Self. Given an integer array nums, return an array counts where counts[i] is the number of elements to the right of nums[i] that are strictly smaller than nums[i].

Constraints

  • 1 ≤ nums.length ≤ 10⁵
  • −10⁴ ≤ nums[i] ≤ 10⁴ (LC); treat keys as unbounded — the treap doesn’t care.

Clarifying Questions

  1. Strictly smaller, or ≤? (Strictly smaller — equal elements don’t count.)
  2. Duplicates in the input? (Yes; the structure must handle equal keys.)
  3. Is the value range bounded? (On LC yes — which enables a BIT over compressed values, lab-03; a treap needs no compression.)
  4. Online or offline queries? (Offline here, but the treap solution is fully online — relevant follow-up.)

Examples

countSmaller([5, 2, 6, 1]) → [2, 1, 1, 0]
countSmaller([-1])         → [0]
countSmaller([-1, -1])     → [0, 0]
countSmaller([2, 0, 1])    → [2, 0, 0]

Initial Brute Force

For each i, scan j > i and count nums[j] < nums[i]. Two nested loops.

Brute Force Complexity

O(n²) time, O(1) extra space. At n = 10⁵: 10¹⁰ comparisons. TLE.

Optimization Path

Three O(n log n) routes: (a) merge-sort with cross-count — offline only; (b) BIT over compressed values (lab-03) — needs bounded/compressible keys, offline compression pass; (c) balanced BST with subtree sizes — online, unbounded keys, no preprocessing. Sweep right→left; for each element, query rank(x) = count of inserted keys < x, then insert x. The treap is route (c) made writable: both the query and the insert are O(log n) expected.

Final Expected Approach

_state = 0x9E3779B97F4A7C15                     # fixed-seed LCG: deterministic runs
def _pri():
    global _state
    _state = (_state * 6364136223846793005 + 1442695040888963407) & (2**64 - 1)
    return _state

class Node:
    __slots__ = ("key", "pri", "size", "l", "r")
    def __init__(self, key):
        self.key, self.pri, self.size = key, _pri(), 1
        self.l = self.r = None

def sz(t):
    return t.size if t else 0

def pull(t):
    t.size = 1 + sz(t.l) + sz(t.r)

def split(t, key):
    """(left, right): left holds keys < key, right holds keys ≥ key."""
    if t is None:
        return None, None
    if t.key < key:
        a, b = split(t.r, key)
        t.r = a
        pull(t)
        return t, b
    a, b = split(t.l, key)
    t.l = b
    pull(t)
    return a, t

def merge(a, b):
    """Every key in a ≤ every key in b (caller's contract)."""
    if a is None: return b
    if b is None: return a
    if a.pri > b.pri:
        a.r = merge(a.r, b)
        pull(a)
        return a
    b.l = merge(a, b.l)
    pull(b)
    return b

def insert(t, key):                              # duplicates allowed
    a, b = split(t, key)
    return merge(merge(a, Node(key)), b)

def erase(t, key):                               # remove one occurrence
    if t is None:
        return None
    if t.key == key:
        return merge(t.l, t.r)
    if key < t.key:
        t.l = erase(t.l, key)
    else:
        t.r = erase(t.r, key)
    pull(t)
    return t

def kth(t, k):                                   # 0-indexed k-th smallest
    while t:
        ls = sz(t.l)
        if k < ls:
            t = t.l
        elif k == ls:
            return t.key
        else:
            k -= ls + 1
            t = t.r
    raise IndexError("k out of range")

def rank(t, key):                                # = count of keys < key
    c = 0
    while t:
        if t.key < key:
            c += sz(t.l) + 1
            t = t.r
        else:
            t = t.l
    return c

def countSmaller(nums):                          # LC 315
    root, out = None, []
    for x in reversed(nums):
        out.append(rank(root, x))
        root = insert(root, x)
    out.reverse()
    return out

Recursion in split/merge/erase runs to tree depth — O(log n) expected (measured: depth 39 after 10⁵ sorted inserts), safely under Python’s default limit of 1000. kth/rank are iterative anyway.

Data Structures Used

  • Node with key, pri, size, l, r and __slots__ (halves memory, speeds attribute access).
  • Fixed-seed 64-bit LCG for priorities — deterministic, reproducible failures.
  • No parent pointers: split/merge never need them.

Correctness Argument

Shape/depth: the treap’s shape equals the BST obtained by inserting keys in decreasing-priority order; random priorities ⇒ uniformly random insertion order. Node x is an ancestor of y iff x has the max priority among the keys between them — probability 2/(d+1) for keys d apart. Summing harmonics gives E[depth] = O(log n), the same recursion-tree analysis as randomized quicksort.

split: induction on tree size. The root routes by BST order into one output; the recursive call splits the relevant subtree; reattaching the returned piece preserves BST order (all its keys are on the correct side of the root) and heap order (parent priorities unchanged). pull on unwind fixes sizes.

merge: precondition max(a) ≤ min(b). The higher-priority root becomes the merged root, preserving heap order; BST order holds since everything in b exceeds everything in a.

rank: walking down, every time we go right past a node with key < target, its left subtree and itself are all < target; nothing else < target is skipped. Sizes are exact because every mutation is followed by pull.

Complexity

OperationTime (expected)Space
split / merge / insert / eraseO(log n)O(n) total
kth / rankO(log n)O(1)
LC 315 totalO(n log n)O(n)

Implementation Requirements

  • pull() after every child reassignment in split/merge/erase — no exceptions.
  • Split convention < key | ≥ key must match rank (strict <) so duplicates aren’t counted as smaller.
  • erase removes exactly one occurrence; kth is 0-indexed — document both.
  • __slots__ on the node class; pure-Python treaps are constant-heavy without it.

Tests

  • The four fixed LC 315 cases above.
  • 300 random arrays (length ≤ 60, values in [−20, 20], duplicates) vs the O(n²) brute force — exact match.
  • 200 random op sequences (insert/erase/kth/rank, 200 ops each) vs a bisect-maintained sorted list oracle, asserting sizes after every op.
  • Adversarial order: 10⁵ sorted inserts; assert depth < 3 log₂ n (measured 39 vs 16.6 log₂).

Follow-up Questions

  • Why do random priorities give E[depth] = O(log n)? → Shape = BST built in decreasing-priority (i.e., random) insertion order. x is y’s ancestor iff x’s priority is max over the key-interval between them: P = 2/(d+1). E[depth] = Σ 2/(d+1) ≈ 2 ln n. Identical to randomized quicksort’s recursion tree — the root is a random pivot for its key range.
  • Range reverse on an array in O(log n)? → Implicit treap: split by size into [0,l), [l,r], (r,…]; set a lazy flip flag on the middle; push down by swapping children and propagating flags. This rope structure is what text editors use for million-line buffers.
  • Treap vs AVL/red-black in an interview? → Same O(log n) asymptotics (treap expected, others worst-case). AVL/RB need rotation case analysis nobody writes correctly in 30 minutes; a treap is two primitives plus composition. Say this trade-off out loud.
  • Make it persistent? → Path-copying: split/merge only touch O(log n) nodes per op; copy them instead of mutating. Old roots remain valid versions — a persistent ordered set/rope for cheap, same trick as persistent segment trees.
  • Skip list — the other writable balanced structure? → Probabilistic level towers, expected O(log n), no rotations at all. Redis chose skip lists for sorted sets: simpler code, cache-friendly forward scans for range queries. Treap wins when you need split/merge as first-class ops.
  • When is the BIT (lab-03) better for LC 315? → Bounded or compressible values, offline: BIT is ~5 lines, tiny constants. The treap wins online, with unbounded keys, or when you also need k-th/delete.

Product Extension

A live leaderboard service: rank(score) answers “what percentile am I?”, kth(k) answers “who’s #k?”, both on a stream of score updates (erase old, insert new). This is exactly Redis ZRANK/ZRANGE — Redis backs it with a skip list; a treap is the interview-writable equivalent. Implicit-treap ropes back collaborative editor buffers: cut/paste of a 10⁶-char range in O(log n).

Language/Runtime Follow-ups

  • Python: nothing built-in. sortedcontainers.SortedList (block-list, third-party) is faster in practice; a treap is the from-scratch answer. __slots__ mandatory; consider array-based nodes (parallel lists) to dodge object overhead.
  • C++: std::set lacks order statistics; __gnu_pbds::tree with tree_order_statistics_node_update gives find_by_order/order_of_key. Treap still needed for split/merge/implicit tricks.
  • Java: TreeMap has no rank; augmenting requires a hand-rolled tree anyway — treap it.
  • Go: no ordered set in the stdlib; treaps are the common CP choice.
  • JS/TS: nothing; for small n a sorted array with binary-search insert (O(n) shift) often passes — know when the constant beats the asymptotics.

Common Bugs

  1. Missing pull(): a child reassignment without a size recompute rots size silently; kth/rank return garbage much later. Every mutation path must end in pull.
  2. Split boundary mismatch: split sending ≤ key left while rank counts < key double-counts duplicates — LC 315 fails only on inputs with equal elements.
  3. Heap direction inconsistency: comparing a.pri > b.pri in one merge branch but effectively min-heap in insert logic — the tree still “works” but degenerates to O(n) depth.
  4. Merging unordered trees: merge(a, b) silently produces a non-BST if the key precondition is violated (e.g., swapped return values from split). Assert the contract in debug builds.
  5. Weak priorities: using id(node) or a constant seed per-node — correlated priorities recreate the sorted-insert worst case. One global LCG or random.random().
  6. kth off-by-one: mixing 0- and 1-indexed k between kth and callers; the k == ls branch is the giveaway to re-check.

Debugging Strategy

Write a debug-only O(n) validator: in-order traversal asserts keys non-decreasing, every parent priority ≥ children, and size == 1 + sz(l) + sz(r) at every node. Run it after every operation in the randomized test; the first failing op localizes the bug to one function. Then shrink: rerun the failing seed with shorter sequences until minimal, and hand-trace split/merge on the 4–5 node result. For LC 315 disagreements, print the in-order key list next to the oracle’s sorted list — a duplicate-boundary bug shows up as an off-by-one exactly at repeated values.

Mastery Criteria

  • Wrote split + merge from memory, correct first run, in <10 minutes.
  • Full ordered set (insert/erase/kth/rank) passing the randomized oracle in <25 minutes.
  • Solved LC 315 three ways from cold start: merge-sort count, BIT, treap.
  • Explained E[depth] = O(log n) via the quicksort correspondence in <2 minutes.
  • Sketched implicit-treap range-reverse (size-split + lazy flip) on a whiteboard.
  • Stated the treap-vs-AVL/RB and treap-vs-skip-list trade-offs unprompted.