Coding interview: How would you use a wavelet matrix for range k-th queries?

Prompt and scope

Given a static integer array a, answer many half-open range queries [l, r): return the k-th smallest value, the frequency of x, and the number of elements in [lo, hi). The array never changes and values may be large. Design and analyze a structure faster than sorting every range.

A wavelet matrix stably partitions values by bits from most significant to least significant, storing a bitvector and prefix-one counts at each level. It needs no explicit tree pointers; each query maps its interval to the next level. State that k is zero-based, handle coordinate compression and duplicates, and give the bounds.

What the interviewer evaluates

• Explain why stable partitioning, the zero-block start, and rank mapping preserve order.
• Select the k-th value over [l, r) and accumulate bits correctly.
• Handle duplicates, empty ranges, out-of-range k, and signed values.
• Distinguish paths for value-domain counts, point frequency, and k-th queries.
• Give O(B) query, O(nB) construction, and compressible-space bounds.
• Recognize that static structure does not provide cheap updates and know alternatives.

Clarifying questions

1. Is r exclusive, and is k zero-based or one-based?
2. Is the array truly immutable? If not, what are update and query rates?
3. Are values signed, and what is their maximum width? Can we coordinate-compress them?
4. What memory is available for rank, and may bitvectors be blocked or compressed?
5. Do we need only k-th, or also frequency, predecessor, or range sum? The operations affect the choice.

30-second answer

I would coordinate-compress values to non-negative codes with bit width B. During construction, stably partition the current sequence from the highest bit down, storing each level's bitvector and prefix rank-one counts. For k-th, keep [l,r), count zeros at the level, and either map to the zero block or subtract zeros and map to the one block while setting that answer bit. Frequency uses two rank walks; a value-domain count is the difference of two countLess calls. Queries cost O(B) and construction costs O(nB).

Step-by-step solution

1. Encode the value domain

For arbitrary signed integers, sort distinct values and map them to 0..m-1, retaining a code-to-value array. Then B is ceil(log2(m)), with an explicit case for m=1. If natural ordering must be preserved directly, flip the sign bit before treating signed values as unsigned.

2. Build one stable level

Inspect cur at bit, append all zero-bit values to next, then all one-bit values, preserving order within both groups. bv[i] records the bit at original position i, and zeroCount is the number of zeros. Stability keeps later intervals tied to the same original elements.

rank1(i) = number of ones in bv[0..i)
zeroCount = n - rank1(n)
for interval [l, r):
  zero interval = [l - rank1(l), r - rank1(r))
  one interval  = [zeroCount + rank1(l), zeroCount + rank1(r))

3. Query the range k-th value

At each level compute zeros = (r-l) - (rank1(r)-rank1(l)). When k is smaller than zeros, map to the zero interval. Otherwise subtract zeros, map to the one interval, and set the current answer bit. After B levels, decode the code back to its original value.

4. Query one-value frequency

Treat each target bit as a fixed branch and map [l,r) the same way. A zero target follows the zero interval; a one target follows the one interval using zeroCount. After B levels, the interval length is the frequency. A target absent from the compressed dictionary returns zero.

5. Query a value range

Define countLess(x, l, r) as the number of values below x in [l,r). At a level where x has bit one, every zero branch is smaller, so add zeros and continue into the one branch. For bit zero, continue only into the zero branch. The count in [lo, hi) is countLess(hi)-countLess(lo).

6. Bounds and verification

Define behavior for an empty range or when the left endpoint is not smaller than the right endpoint; never index rank arrays outside their bounds. Require k to fall within the current range length. Test all-equal, sorted, interleaved duplicates, negative values, one-element ranges, maximum bit width, and values absent from the dictionary, comparing each result with a brute-force sort or count.

7. Complexity and trade-offs

With ordinary prefix counts, each level stores O(n) counters, so space and construction are O(nB) and every operation is O(B). A rank-supporting compressed bitvector reduces space and constants. The structure fits immutable, query-heavy workloads. For updates, consider blocked rebuilds, dynamic bitvectors, a segment tree of ordered sets, or offline processing and reevaluate memory and update costs.

Sample strong answer

I would state that ranges are half-open, k is zero-based, and the array is immutable. I would coordinate-compress values and use B bits. Construction stably partitions from the highest bit down, retaining prefix rank-one counts and zero-block lengths at every level.

For k-th, each level counts zeros in the current interval. If k belongs to zeros, map with l-rank1(l) and r-rank1(r); otherwise subtract zeros, map with zeroCount+rank1(l) and zeroCount+rank1(r), and set the answer bit. Frequency follows a fixed value path, while value-domain count is two countLess calls. Construction is O(nB) and each query O(B), with explicit errors or zero for invalid ranges and absent codes.

Common mistakes

• Mixing closed and half-open ranges → rank shifts by one → use [l,r) consistently and write the mapping.
• Forgetting stable partition → later intervals no longer identify the same elements → preserve order in both blocks.
• Entering the one block without subtracting zeros → k-th values are too large → subtract before mapping.
• Comparing signed values as unsigned bits → negatives are misordered → compress or flip the sign bit.
• Assuming updates are cheap → updates invalidate level permutations → state the static precondition and alternatives.
• Testing only distinct values → duplicates and bounds bugs remain hidden → test equal, interleaved, empty, and invalid cases.

Follow-up questions and responses

Why not sort every range?

Sorting one range costs O((r-l) log(r-l)) and repeats work across queries. The matrix precomputes branch information, so a query visits only B levels and suits static, high-query workloads.

Why does rank1 map intervals?

Prefix rank tells how many ones occur before each endpoint, which gives the interval's relative positions in the zero and one blocks. Stable partitioning ensures those positions represent the same elements.

How do you answer k-th largest?

Convert it to k-th smallest with length - 1 - k, or prefer the one branch at every level while subtracting its count. Both remain O(B).

What if the value domain is much larger than n?

Coordinate-compress observed values and retain the reverse map. For an unseen query value, binary-search its insertion boundary or return frequency zero.

What if updates are required?

A plain wavelet matrix is not update-friendly. Use blocked rebuilding, dynamic bitvectors, a segment tree of ordered structures, or offline processing based on update/query ratios, latency goals, and memory.