Stage 4: Advanced Techniques·Segment Tree & BIT

Hard#69910 min readSegment TreeCoordinate CompressionLazy PropagationLeetCode

Falling Squares

Coordinate-compressed segment tree with lazy max-assign

Coordinate compression + lazy propagation. A complete segment tree exercise.

Published Apr 16, 2026

Problem

Each element of positions is a falling square [left, sideLength]. Squares fall one at a time onto the number line; each square stacks on top of whatever intervals it overlaps. After every drop, return the height of the tallest stack on the line so far.

Formally, when square i lands it occupies the interval [left, left + sideLength). Its base height is max(heights of any previously landed square that overlaps that interval). Its top height is base + sideLength. The answer at step i is the running max of all top heights observed through square i.

Input

positions = [[1, 2], [2, 3], [6, 1]]

Output

[2, 5, 5]

Explanation

Square 1 lands flat at height 2. Square 2 overlaps it, stacks to height 2+3=5. Square 3 is isolated at height 1, global max stays 5.

Input

positions = [[100, 100], [200, 100]]

Output

[100, 100]

Explanation

Squares are edge-adjacent, not overlapping (intervals [100,200) and [200,300) share only an endpoint). Both land at height 100.

Intuition

Think of the number line as a landscape. A falling square asks two things of the landscape under its footprint:

Range max. What's the tallest point currently underneath?
Range assignment. Set every point under my footprint to my new top height.

Range max + range assign is the canonical segment tree with lazy propagation workload. The only twist: coordinates can be up to 10^8, far too large for a raw tree. We compress the set of distinct endpoints that actually appear in the input — at most 2n of them — and the tree shrinks to O(n) leaves.

The interval for square [l, s] is [l, l+s). We represent each half-open piece by a contiguous band of compressed leaves. Two squares share a leaf exactly when their footprints overlap.

Approach

Brute Force per Square

TimeO(n²)

SpaceO(n)

Keep every landed square. For each new square, scan the list and find the max height of any previously landed square whose footprint overlaps. Place this square at that_max + side. Track the running global max.

impl Solution {
    pub fn falling_squares(positions: Vec<Vec<i32>>) -> Vec<i32> {
        let mut landed: Vec<(i32, i32, i32)> = Vec::new(); // (l, r, top)
        let mut out: Vec<i32> = Vec::with_capacity(positions.len());
        let mut best = 0i32;
        for p in &positions {
            let (l, side) = (p[0], p[1]);
            let r = l + side;
            let mut base = 0i32;
            for &(ll, rr, t) in &landed {
                // Half-open overlap: [l, r) ∩ [ll, rr) non-empty.
                if l < rr && ll < r {
                    if t > base {
                        base = t;
                    }
                }
            }
            let top = base + side;
            landed.push((l, r, top));
            if top > best {
                best = top;
            }
            out.push(best);
        }
        out
    }
}

Rust · runnable

function fallingSquares(positions: number[][]): number[] {
  const landed: Array<[number, number, number]> = []; // [l, r, top]
  const out: number[] = [];
  let best = 0;
  for (const [l, side] of positions) {
    const r = l + side;
    let base = 0;
    for (const [ll, rr, t] of landed) {
      if (l < rr && ll < r && t > base) base = t;
    }
    const top = base + side;
    landed.push([l, r, top]);
    if (top > best) best = top;
    out.push(best);
  }
  return out;
}

TypeScript · runnable

Fine up to a few thousand squares. Gets quadratic on the real constraints (n ≤ 10^3 at LeetCode, but the pattern scales to much larger inputs at O(n²) — we'll beat it next).

Segment Tree with Lazy Max-Assign

TimeO(n log n)

SpaceO(n)

Recommended

Compress all endpoints l and l + side - 1 from every square into a sorted set, and give each distinct coordinate an index. Each square now maps to a contiguous leaf range in the compressed tree.

Why l + side - 1 instead of l + side? We want squares that only share an endpoint (like [100, 200) and [200, 300)) to not overlap. Using l + side - 1 as the inclusive right endpoint lets us treat the compressed tree as a sequence of closed intervals: two squares share compressed indices exactly when their half-open footprints share a point.

Each tree node stores:

max — the current maximum height over the node's interval.
lazy — a pending "assign this value to every descendant" tag (0 = no tag).

The assignment is a set, not an add — everything in the range becomes value, overwriting whatever stacked before. Lazy propagation pushes the tag down before descending.

impl Solution {
    pub fn falling_squares(positions: Vec<Vec<i32>>) -> Vec<i32> {
        // Coordinate compression over inclusive endpoints l, l+side-1.
        let mut pts: Vec<i32> = Vec::with_capacity(positions.len() * 2);
        for p in &positions {
            pts.push(p[0]);
            pts.push(p[0] + p[1] - 1);
        }
        pts.sort_unstable();
        pts.dedup();
        let m = pts.len();
        let idx = |v: i32| -> usize { pts.binary_search(&v).unwrap() };

        let size = m.max(1);
        let mut seg = SegTree::new(size);
        let mut out: Vec<i32> = Vec::with_capacity(positions.len());
        let mut best = 0i32;
        for p in &positions {
            let l = idx(p[0]);
            let r = idx(p[0] + p[1] - 1);
            let base = seg.query(1, 0, size - 1, l, r);
            let top = base + p[1];
            seg.update(1, 0, size - 1, l, r, top);
            if top > best {
                best = top;
            }
            out.push(best);
        }
        out
    }
}

struct SegTree {
    max: Vec<i32>,
    lazy: Vec<i32>, // 0 means no pending assignment
}

impl SegTree {
    fn new(n: usize) -> Self {
        Self {
            max: vec![0; 4 * n],
            lazy: vec![0; 4 * n],
        }
    }

    fn push_down(&mut self, node: usize) {
        if self.lazy[node] != 0 {
            let v = self.lazy[node];
            self.max[2 * node] = v;
            self.lazy[2 * node] = v;
            self.max[2 * node + 1] = v;
            self.lazy[2 * node + 1] = v;
            self.lazy[node] = 0;
        }
    }

    fn update(&mut self, node: usize, lo: usize, hi: usize, ql: usize, qr: usize, val: i32) {
        if qr < lo || hi < ql {
            return;
        }
        if ql <= lo && hi <= qr {
            self.max[node]= val;
            self.lazy[node]= val;
            return;
        }
        self.push_down(node);
        let mid= (lo + hi) / 2;
        self.update(2 * node, lo, mid, ql, qr, val);
        self.update(2 * node + 1, mid + 1, hi, ql, qr, val);
        self.max[node]= self.max[2 * node].max(self.max[2 * node + 1]);
    }

    fn query(&mut self, node: usize, lo: usize, hi: usize, ql: usize, qr: usize) -> i32 {
        if qr < lo || hi < ql {
            return 0;
        }
        if ql <= lo && hi <= qr {
            return self.max[node];
        }
        self.push_down(node);
        let mid= (lo + hi) / 2;
        let a= self.query(2 * node, lo, mid, ql, qr);
        let b= self.query(2 * node + 1, mid + 1, hi, ql, qr);
        a.max(b)
    }
}

Rust · runnable

function fallingSquares(positions: number[][]): number[] {
  const pts: number[] = [];
  for (const [l, s] of positions) {
    pts.push(l, l + s - 1);
  }
  pts.sort((a, b) => a - b);
  const dedup: number[] = [];
  for (const p of pts) {
    if (dedup.length === 0 || dedup[dedup.length - 1] !== p) dedup.push(p);
  }
  const m = Math.max(dedup.length, 1);
  const idx = (v: number): number => {
    let lo = 0;
    let hi = dedup.length - 1;
    while (lo <= hi) {
      const mid = (lo + hi) >> 1;
      if (dedup[mid] === v) return mid;
      if (dedup[mid] < v) lo = mid + 1;
      else hi = mid - 1;
    }
    return lo;
  };

  const max = new Array<number>(4 * m).fill(0);
  const lazy = new Array<number>(4 * m).fill(0);

  const push = (node: number): void => {
    if (lazy[node] !== 0) {
      const v = lazy[node];
      max[2 * node] = v;
      lazy[2 * node] = v;
      max[2 * node + 1] = v;
      lazy[2 * node + 1] = v;
      lazy[node] = 0;
    }
  };

  const update = (
    node: number,
    lo: number,
    hi: number,
    ql: number,
    qr: number,
    val: number,
  ): void => {
    if (qr < lo || hi < ql) return;
    if (ql <= lo && hi <= qr) {
      max[node]= val;
      lazy[node]= val;
      return;
    }
    push(node);
    const mid= (lo + hi) >> 1;
    update(2 * node, lo, mid, ql, qr, val);
    update(2 * node + 1, mid + 1, hi, ql, qr, val);
    max[node] = Math.max(max[2 * node], max[2 * node + 1]);
  };

  const query = (
    node: number,
    lo: number,
    hi: number,
    ql: number,
    qr: number,
  ): number => {
    if (qr < lo || hi < ql) return 0;
    if (ql <= lo && hi <= qr) return max[node];
    push(node);
    const mid= (lo + hi) >> 1;
    return Math.max(
      query(2 * node, lo, mid, ql, qr),
      query(2 * node + 1, mid + 1, hi, ql, qr),
    );
  };

  const out: number[] = [];
  let best = 0;
  for (const [l, s] of positions) {
    const lo = idx(l);
    const hi = idx(l + s - 1);
    const base = query(1, 0, m - 1, lo, hi);
    const top = base + s;
    update(1, 0, m - 1, lo, hi, top);
    if (top > best) best = top;
    out.push(best);
  }
  return out;
}

TypeScript · runnable

The key invariant for assign lazy: when a tag is set, every child's max and lazy are overwritten on the next descent. Earlier partial updates are dropped — exactly the behavior a falling square wants.

Trace

positions = [[1, 2], [2, 3], [6, 1]]. Endpoints are {1, 2, 2, 4, 6, 6} → compressed to {1, 2, 4, 6} with indices 0..3.

step	l..r (compressed)	query max	top	seg max after	running best
1	0..1 (vals 1..2)	0	0 + 2 = 2	[2, 2, 0, 0]	2
2	1..2 (vals 2..4)	2	2 + 3 = 5	[2, 5, 5, 0]	5
3	3..3 (val 6)	0	0 + 1 = 1	[2, 5, 5, 1]	5

Square 1 claims compressed leaves {1, 2} (raw endpoints 1 and 2). Square 2's footprint is [2, 5) → compressed {2, 4} → leaves 1..2, overlaps square 1 at leaf 1, stacks on top. Square 3's footprint is [6, 7) → leaf 3, isolated.

Edge cases

Edge-adjacent squares don't overlap. [[100, 100], [200, 100]] land at height 100 each. Mapping to inclusive l + side - 1 is the coordinate trick that preserves this.
Full overlap. A new square fully inside a taller one still ends up at the taller one's top plus its own side — the query returns the max, not the average.
Single square. Compressed array has two distinct points (or one, if side == 1); the tree size floors at 1 leaf.
Large coordinates. Raw left + side can be near 2 * 10^8; compression keeps the tree size at O(n).

💡

Tip

assign lazy is subtler than add lazy — the tag completely overwrites. If you test with add semantics by accident, you'll double-count overlapping drops.

Takeaway

Range max + range assign + coordinate compression is the template for every sweep-style problem where intervals paint over each other. Ad-hoc "who's taller right now" in a live order book, running dashboards of top-of-book depth, and any "game state over a timeline" reduce to the same three primitives.

Once the lazy-assign skeleton is cached in your fingers, #2407 is a shorter version of the same trick with max updates only and no assignments.

Medium#307

Range Sum Query - Mutable

Classic segment tree. Build, update, query. Implement from scratch in Rust.

Hard#315

Count of Smaller Numbers After Self

BIT or merge sort. Process right to left, count elements smaller than current.

Hard#2407

Longest Increasing Subsequence II

Segment tree optimization on DP transitions. The ‘DP + data structure’ pattern.