公式

A - セグメント木/Segment Tree 解説 by maspy

[1] セグメント木とは

セグメント木は，列に対する区間クエリを高速に処理するためのデータ構造です．今回は，\(1\) 点変更と区間最小値取得をそれぞれ \(O(\log N)\) 時間で処理します．

セグメント木では，完全二分木の各ノードに，ある区間の情報を保持します．今回の場合，各ノードには対応する区間の最小値を保持します．

重要なのは，次の 不変条件 です．

  • 葉には，列 \(A\) の各要素を保持する．
  • 葉以外のノードには，\(2\) つの子が保持する値の最小値を保持する．

不変条件とは，データ構造が常に満たすべき性質のことです．セグメント木に限らず，データ構造の実装方法を考えたり，正当性を証明したりする際には，不変条件を意識することが重要です．

次節で説明するセグメント木の実装でも，この不変条件を意識するようにしてください．

[2] セグメント木の実装

[2.1] セグメント木のインデックス

セグメント木の実装では，完全二分木の各ノードに整数の番号を付けます．

  • 根の番号は \(1\)．
  • ノード \(i\) の子の番号は \(2i,2i+1\)．
  • ノード \(i\) の親の番号は \(\lfloor i/2\rfloor\)．

一般に木構造を実装する際には，ポインタなどを利用して親子関係を管理することがあります．一方，セグメント木では，ノードの番号から親や子の番号を簡単に計算できます．そのため，配列だけで実装でき，コードが簡潔で高速になるという利点があります．

考える列の長さが \(N\) である場合，完全二分木の葉の個数を \(N\) 以上の最小の \(2\) 冪 seg_size に設定します．

各ノードの情報を配列 dat に格納することにすると，元の列の要素 A[i] は，dat[seg_size+i] に対応します．

以下では，列の添字を \(0\) から始め，dat，seg_size をあらかじめ定義しているものとします．また，INF は問題で扱うすべての要素の値よりも大きい整数とします．

[2.2] 初期構築と 1 点変更

初期構築では，葉に初期値を代入し，下から順に各ノードの値を計算します．

void build(const vector<int>& A) {
  int N = A.size();
  seg_size = 1;
  while (seg_size < N) seg_size *= 2;

  // 葉への代入
  for (int i = 0; i < seg_size; ++i) {
    dat[seg_size + i] = (i < N ? A[i] : INF);
  }

  // 下から順に各ノードの値を計算
  for (int i = seg_size - 1; i >= 1; --i) {
    dat[i] = min(dat[2 * i], dat[2 * i + 1]);
  }
}

葉の個数を \(2\) 冪にそろえるため，余分な葉には INF を代入しています．

葉は \(O(N)\) 個であり，完全二分木全体のノード数も \(O(N)\) 個なので，初期構築の計算量は \(O(N)\) です．

\(1\) 点変更では，対応する葉の値を書き換えたあと，その葉の祖先を順に再計算します．

void point_set(int i, int x) {
  i += seg_size;
  dat[i] = x;

  while (i > 1) {
    // 親に移動し，不変条件を保つように再計算する
    i /= 2;
    dat[i] = min(dat[2 * i], dat[2 * i + 1]);
  }
}

変更した葉の祖先だけを再計算すれば，[1] の不変条件を保つことができます．

完全二分木の高さは \(O(\log N)\) なので，\(1\) 点変更の計算量も \(O(\log N)\) です．

このように，初期構築と \(1\) 点変更は，親子のインデックスの関係と不変条件を理解していれば，比較的容易に実装できます．

次に，区間最小値取得を考えます．ここでは，半開区間 \([L,R)\) の最小値を求める関数 prod(L,R) を実装します．

主な実装方法として，再帰関数を用いる方法と，用いない方法の \(2\) 通りを紹介します．

[2.3] 区間クエリの再帰的な実装

再帰関数を用いる場合，各ノードが担当する区間を引数で管理します．

ノード \(i\) が担当する区間を \([\mathrm{node\_L},\mathrm{node\_R})\)，求めたい区間を \([L,R)\) とします．これらの区間の関係に応じて，次の \(3\) 通りに場合分けします．

  1. 共通部分がない場合：INF を返す．
  2. \([\mathrm{node\_L},\mathrm{node\_R})\) が \([L,R)\) に含まれる場合：dat[i] を返す．
  3. それ以外の場合：左右の子を再帰的に調べ，結果の最小値を返す．

この方針をそのまま実装すると，次のようになります．

// ノード i は区間 [node_L, node_R) を担当する
// 返り値は区間 [node_L, node_R) と [L, R) の共通部分の最小値
int prod_rec(int i, int node_L, int node_R, int L, int R) {
  // 共通部分がない
  if (node_R <= L || R <= node_L) return INF;

  // ノードの区間全体が含まれる
  if (L <= node_L && node_R <= R) return dat[i];

  // 左右の子を再帰的に処理する
  int m = (node_L + node_R) / 2;
  int x = prod_rec(2 * i, node_L, m, L, R);
  int y = prod_rec(2 * i + 1, m, node_R, L, R);
  return min(x, y);
}

int prod(int L, int R) {
  return prod_rec(1, 0, seg_size, L, R);
}

この実装では，求めたい区間がノードの担当区間を完全に含んでいる場合，そのノードの値を利用して，それ以上の探索を行いません．

つまり，求めたい区間をいくつかのノードが表す区間に分解し，それらの最小値を求めています．

半開区間 \([L,R)\) の両端の境界をまたぐノードだけが再帰的に分割されます．深さが一定のところではそのようなノードは高々 \(2\) 個なので，この実装の計算量は \(O(\log N)\) です．

[2.4] 区間クエリの非再帰的な実装

再帰関数を用いず，区間の左右端から処理する方法もあります．

まず，L と R に seg_size を加えて葉の番号に変換します．その後，区間の両端から必要なノードを選びながら，親の階層へ移動します．

int prod(int L, int R) {
  L += seg_size;
  R += seg_size;

  int ANS = INF;
  while (L < R) {
    if (L & 1) ANS = min(ANS, dat[L++]);
    if (R & 1) ANS = min(ANS, dat[--R]);
    L /= 2;
    R /= 2;
  }
  return ANS;
}

この実装では，各ループの開始時点で，未処理の区間がちょうどある深さのノード番号 \([L,R)\) に対応するようにしています．

ある深さのノードについて，番号 \(2i,2i+1\) のノードが両方とも計算対象となっている場合，それらを共通の親 \(i\) にまとめることができます．このような \(2\) ノードずつのペアを作っていったときに，左端と右端にペアを作れず余るノードができる場合があります．

  • L が奇数の場合：最も左のノード L はペアを作れません．このノードの値を答えに反映し，L を \(1\) 増やします．
  • R が奇数の場合：最も右のノード R-1 はペアを作れません．このノードの値を答えに反映し，R を \(1\) 減らします．

これらの処理のあと，残ったノードはすべて兄弟同士のペアになります．各ペアを親ノードに置き換えることで，ひとつ上の階層ではノード番号 \([\lfloor L/2\rfloor,\lfloor R/2\rfloor)\) が計算対象になります．

これを L == R になるまで繰り返すことで，求めたい区間をいくつかのノードが表す区間に分解できます．

各階層で高々 \(2\) 個のノードの値を答えに反映し，階層数は \(O(\log N)\) なので，計算量は \(O(\log N)\) です．

再帰的な実装と比較すると，少し独特なインデックス操作が必要ですが，コードが短く，実行速度の面でも有利です．

[2.5] 解答例

以上の実装をまとめた解答例を次に示します．

セグメント木の実装は初心者向けの代表的な学習テーマであり，インターネット上にも解説が多数あります．必要に応じて外部資料も参照してください．例えば，私の記事

やその参考文献をあたるとよいと思います．

[3] AtCoder Library を使う場合

AtCoder Library には，セグメント木の実装 atcoder::segtree が用意されています．

利用する際には，管理するデータの型，\(2\) つの区間情報を結合する演算 op，およびその単位元 e を指定します．

今回の場合は，管理するデータの型を int，演算を min，単位元を十分大きな整数とすればよいです．例えば，次のように定義できます．

const int INF = 1'000'000'001;

int op(int a, int b) {
  return min(a, b);
}

int e() {
  return INF;
}

これらを指定することで，初期列 A からセグメント木を構築できます．

atcoder::segtree<int, op, e> seg(A);

構築後は，次の操作を使用します．

  • seg.set(i,x)：位置 \(i\) の要素を \(x\) に変更する．
  • seg.prod(L,R)：半開区間 \([L,R)\) の最小値を求める．

いずれも \(O(\log N)\) 時間で処理できます．

[4] セグメント木の利用方法の学習について

今回扱った区間最小値のほかにも，セグメント木では，区間和・区間最大値・行列積など，さまざまな区間クエリを扱うことができます．

どのような演算をセグメント木で扱えるのかを理解するうえでは，モノイドという代数構造が重要です．AtCoder Library の segtree も，モノイドの考え方に基づいて実装されています．モノイドの定義や，簡単な具体例については次の講座を参照してください．

より高度なモノイドを自力で考え，区間クエリに適用するための考え方については次の講座を参照してください．

また，本コンテストはセグメント木の学習や，理解度の確認にもご利用いただけます．