TypeScriptで遅延セグメント木を実装する：区間加算と区間最大値を不変条件から理解する

配列の連続区間に値を加えながら、その区間の最大値を何度も求める問題を考えます。素直に配列を書き換えると、区間更新も区間最大値も最悪で O(n) です。遅延セグメント木なら、どちらも O(log n) にできます。

この記事では、区間を半開区間 [left, right) で統一し、次の二つだけを実装します。

rangeAdd(left, right, delta): 区間の全要素に delta を加える
rangeMax(left, right): 区間の最大値を返す

遅延セグメント木は実装量が多く見えますが、守るべき不変条件は二つです。

max[node] は、そのノードの担当区間に加えられた更新をすべて反映した最大値です。
lazy[node] は、子孫にはまだ渡していない加算値です。ただし max[node] には反映済みです。

この二つを混同しなければ、更新と問い合わせを同じ再帰の形で書けます。

なぜ通常のセグメント木だけでは足りないのか

通常のセグメント木では一点更新と区間問い合わせを O(log n) で扱えます。しかし [100, 900) のような広い区間に +5 を加えるたび、葉まで降りると O(n) に戻ってしまいます。

遅延伝播では、担当区間が更新範囲に完全に含まれるノードで止まり、そこにだけ +5 を記録します。子の値が必要になったときだけ、その記録を子へ渡します。この「必要になるまで降ろさない」ことが lazy propagation です。

実装

配列サイズを 2 のべき乗に丸める必要はありません。再帰呼び出しで担当区間 [nodeLeft, nodeRight) を明示すれば、元の長さのまま扱えます。

class LazyMaxSegmentTree {
  readonly #size: number;
  readonly #max: Float64Array;
  readonly #lazy: Float64Array;

  constructor(values: readonly number[]) {
    if (values.length === 0) {
      throw new RangeError("values must not be empty");
    }

    this.#size = values.length;
    this.#max = new Float64Array(values.length * 4);
    this.#lazy = new Float64Array(values.length * 4);
    this.#build(1, 0, this.#size, values);
  }

  rangeAdd(left: number, right: number, delta: number): void {
    this.#assertRange(left, right);
    this.#add(1, 0, this.#size, left, right, delta);
  }

  rangeMax(left: number, right: number): number {
    this.#assertRange(left, right);
    return this.#query(1, 0, this.#size, left, right);
  }

  #build(
    node: number,
    nodeLeft: number,
    nodeRight: number,
    values: readonly number[],
  ): void {
    if (nodeRight - nodeLeft === 1) {
      this.#max[node] = values[nodeLeft];
      return;
    }

    const middle = Math.floor((nodeLeft + nodeRight) / 2);
    this.#build(node * 2, nodeLeft, middle, values);
    this.#build(node * 2 + 1, middle, nodeRight, values);
    this.#pull(node);
  }

  #apply(node: number, delta: number): void {
    this.#max[node] += delta;
    this.#lazy[node] += delta;
  }

  #push(node: number, nodeLeft: number, nodeRight: number): void {
    if (nodeRight - nodeLeft === 1 || this.#lazy[node] === 0) {
      return;
    }

    this.#apply(node * 2, this.#lazy[node]);
    this.#apply(node * 2 + 1, this.#lazy[node]);
    this.#lazy[node] = 0;
  }

  #pull(node: number): void {
    this.#max[node] = Math.max(this.#max[node * 2], this.#max[node * 2 + 1]);
  }

  #add(
    node: number,
    nodeLeft: number,
    nodeRight: number,
    left: number,
    right: number,
    delta: number,
  ): void {
    if (right <= nodeLeft || nodeRight <= left) {
      return;
    }

    if (left <= nodeLeft && nodeRight <= right) {
      this.#apply(node, delta);
      return;
    }

    this.#push(node, nodeLeft, nodeRight);
    const middle = Math.floor((nodeLeft + nodeRight) / 2);
    this.#add(node * 2, nodeLeft, middle, left, right, delta);
    this.#add(node * 2 + 1, middle, nodeRight, left, right, delta);
    this.#pull(node);
  }

  #query(
    node: number,
    nodeLeft: number,
    nodeRight: number,
    left: number,
    right: number,
  ): number {
    if (right <= nodeLeft || nodeRight <= left) {
      return Number.NEGATIVE_INFINITY;
    }

    if (left <= nodeLeft && nodeRight <= right) {
      return this.#max[node];
    }

    this.#push(node, nodeLeft, nodeRight);
    const middle = Math.floor((nodeLeft + nodeRight) / 2);
    return Math.max(
      this.#query(node * 2, nodeLeft, middle, left, right),
      this.#query(node * 2 + 1, middle, nodeRight, left, right),
    );
  }

  #assertRange(left: number, right: number): void {
    if (
      !Number.isInteger(left) ||
      !Number.isInteger(right) ||
      left < 0 ||
      left >= right ||
      right > this.#size
    ) {
      throw new RangeError("invalid half-open range");
    }
  }
}

`apply` では max と lazy を同時に更新する

完全に覆われた区間では、子へは降りません。それでも親区間の最大値を次の問い合わせで即座に返せなければならないため、#apply は #max と #lazy の両方を更新します。

ここで lazy だけを増やすと、完全一致の rangeMax が古い最大値を返します。逆に max だけを増やすと、部分区間へ次に降りたとき、子が親の更新を知らないままになります。

`push` は部分区間へ降りる直前だけに呼ぶ

#add と #query が担当区間を完全に使えるときは、そのノードで返すか更新して終わります。部分的に重なるときだけ #push が必要です。

この順序には理由があります。親の lazy を先に子へ反映しておけば、子を更新してから #pull した値が、親の不変条件を再び満たします。

小さな例

初期配列を [1, 3, -2, 4] とします。

rangeAdd(1, 4, 5) の後は [1, 8, 3, 9] です。
rangeMax(0, 4) は 9 を返します。
rangeAdd(0, 2, -4) の後は [-3, 4, 3, 9] です。
rangeMax(0, 3) は 4 を返します。

右端を含めない半開区間にしているため、長さは常に right - left で表せます。隣接する [0, 2) と [2, 4) も重ならず、再帰の境界条件をそのまま書けます。

配列とのランダム差分テスト

複雑な木構造は、例だけでは境界バグを見逃しがちです。小さい配列を正解オラクルとして持ち、同じ更新を繰り返して問い合わせ結果を比較します。

import { strict as assert } from "node:assert";

const tree = new LazyMaxSegmentTree([1, 3, -2, 4]);
tree.rangeAdd(1, 4, 5);
assert.equal(tree.rangeMax(0, 4), 9);
tree.rangeAdd(0, 2, -4);
assert.equal(tree.rangeMax(0, 3), 4);
assert.throws(() => tree.rangeAdd(2, 2, 1), RangeError);

let seed = 0x12345678;
const next = () => {
  seed = (seed * 1664525 + 1013904223) >>> 0;
  return seed;
};

const actual = Array.from({ length: 23 }, (_, index) => index - 11);
const expected = [...actual];
const randomTree = new LazyMaxSegmentTree(actual);

for (let step = 0; step < 500; step += 1) {
  const left = next() % expected.length;
  const right = left + 1 + (next() % (expected.length - left));

  if (next() % 2 === 0) {
    const delta = (next() % 21) - 10;
    randomTree.rangeAdd(left, right, delta);
    for (let index = left; index < right; index += 1) {
      expected[index] += delta;
    }
  } else {
    assert.equal(
      randomTree.rangeMax(left, right),
      Math.max(...expected.slice(left, right)),
    );
  }
}

console.log("all assertions passed");

このテストは、重なり方が異なる 500 回の操作で、遅延セグメント木と素直な配列の結果を比較します。特に「更新後に別の部分区間を問い合わせる」ケースで push 忘れを検出できます。

面接で説明するなら

説明は次の順で組み立てると、実装の意図が伝わります。

通常のセグメント木では区間更新が O(n) になるため、完全に覆われたノードに更新を保留する。
max は常に最新、lazy は子への未伝播分という二つの不変条件を置く。
部分区間へ降りる直前だけ push し、戻るときに pull する。
各操作は木の高さに沿って進むため O(log n) になる。
半開区間と差分テストで、端点と伝播順のバグを防ぐ。

遅延セグメント木の難しさは再帰そのものではなく、「今のノードは正しいが、子にはまだ伝えていない」という状態を明確に保つことです。まず区間演算を一種類に絞って不変条件をテストし、必要になってから区間和や代入更新へ拡張するのが安全です。