この記事の概要
セグメント木とは
セグメント木の基本構造
セグメント木の初期実装
- ソースコード例
値の取得
- ソースコード例
値の更新
- ソースコード例
使用例
あとがき
参考にさせていただいた記事

この記事の概要

競技プログラミングの問題を解くにあたって、セグメント木 (Segment tree) がかなり有用で使う場面もそこそこあるな、と感じたので個人的にまとめてみます。

徐々に加筆していくと思います。

セグメント木とは

競技プログラミングで用いるセグメント木とは、主に完全二分木によって実装されたデータ構造で、区間内での値を取得することに優れたものです。

（ここで、「主に」と書いたのは、競プロ界隈では完全二分木によるものを指しますが、それ以外の界隈では必ずしも完全二分木ではない *1 みたいだからです。）

この記事では、完全二分木によるものを扱います。すなわち、整数 $m$ を用いて、要素数が ${2}^{m} - 1$ 個である二分木を用いてセグメント木を表現します。

セグメント木の基本構造

さて、要素数が ${2}^{m} - 1$ である完全二分木は以下のように $m$ 段の構造をしています。

この二分木に対して、上の段からインデックスを振り分けると以下のようになります。

ここで、整数 $n$ を $n = {2}^{m - 1}$ として定めます。以上のように作られた完全二分木に対して、以下のように $0, 1, \cdots, n - 1$ の値に対しての範囲を振り分けます。

このように振り分けられた範囲に対しての値を、それぞれの二分木の頂点に入れることによって、セグメント木が完成します。

例えば、数列 ${a}_{0}, {a}_{1}, \cdots , {a}_{n - 1}$ について、ある区間の最小値を求めるときについては、

頂点 $0$ には ${a}_{0}, {a}_{1}, \cdots , {a}_{n - 1}$ の最小値を入れる
頂点 $1$ には ${a}_{0}, {a}_{1}, \cdots , {a}_{(n / 2) - 1}$ の最小値を入れる
頂点 $2$ には ${a}_{n / 2} , \cdots , {a}_{n - 1}$ の最小値を入れる
$\vdots$
頂点 ${2}^{n} - 2$ には ${a}_{n - 1}$ の最小値（すなわち ${a}_{n - 1}$ ）を入れる

というようにすることで、セグメント木が完成されます。

このようなセグメント木に対して、例えば「 ${a}_{1}, {a}_{2}, \cdots , {a}_{n - 2}$ の最小値を求めたい」というときには、以下の図の色の付いた部分のみを調べ上げて、その中での最小値を導出することで求めることができます。

具体的な値の取得方法は後で詳しく述べますが、このようなセグメント木を用意することで、ある区間での値の取得を $O(\log n)$ の計算量で行うことができます。

セグメント木の初期実装

ここからは、C++でのセグメント木の実装について記します。

まずはセグメント木となる二分木を vector の形で実装します。ここまでの解説の通り、セグメント木は完全二分木で実装されるため、要素数が ${2}^{m} - 1$ の形になります。

すなわち、セグメント木を用いて管理する数列の要素数を $n$ とするとき、

$n \leq {2}^{m - 1}$ となる最小の $m$ の値を求めて、新たに $n = {2}^{m - 1}$ とする
vector の要素数を ${2}^{m} - 1 \, ( = 2n - 1)$ として初期化する

という手順によって、セグメント木の初期化を行えます。

ソースコード例

これをC++にて実装すると、以下のようになります。

struct SegmentTree {
  int n; // セグメント木を用いて管理する数列の要素数
  vector<ll> node; // セグメント木の頂点となる vector

  void init() {
    int sz = 1; // 管理する数列の要素数（2 ^ (m - 1) の形に最終的になる）
    while (sz < n) {
      sz *= 2; // n <= sz となるまで sz の値を2倍する
    }
    n = sz; // n の値を更新する
    node.resize(2 * n - 1); // node の要素数を 2 * n - 1 にする
    return;
  }
}

この init() 関数について、例えば「最小値を求める」のようなものであれば、node の要素数を指定する際に、中身の値としてすべて ${10}^{18}$ のような大きな値を入れることによって、その後の操作を行いやすくなります。

値の取得

セグメント木を用いて、 $[a, \, b)$ の範囲の目的関数の値を求めることについて考えます。

ここで、セグメント木の頂点 $k$ が、 $[l, \, r)$ の範囲を扱っているとします。

範囲 $[a, \, b)$ と範囲 $[ l, \, r )$ の関係次第で、目的関数の値が変わっていきます。どのような関係があるのか考えましょう。

(i) $r \leq a$ または $b \leq l$ のとき

このときは、範囲 $[a, \, b)$ と範囲 $[l, \, r)$ は全く重複していません。

従って、この場合では、目的とするものに影響を与えない値を返すこととなります。これは、例えば目的関数が以下のようなときに

目的関数が最小値のとき: ${10}^{18}$ のような、最大の値を返す
目的関数が最大値のとき: $- {10}^{18}$ のような、最小の値を返す
目的関数が和のとき: $0$ を返す

というようにします。

(ii) $a \leq l$ かつ $r \leq b$ のとき

このときは、範囲 $[a, \, b)$ の中に範囲 $[l, \, r)$ が完全に入っていることとなります。

すなわち、範囲 $[a, \, b)$ の目的関数の値の探索中に範囲 $[l , \,r)$ での探索が入っている場合なので、このときは $\mathrm{node [} k \mathrm{ ]}$ の値を返すというようにします。

(iii) それ以外のとき

このときは、範囲 $[ a, \, b)$ と範囲 $[l, \, r)$ の関係は以下のようになります。

図で表すように、この場合は範囲が入り組んでいるので、インデックス $k$ の更に下のノードについて考える必要があります。ここで、インデックス $k$ の下にあるノードは以下のようになっています。

ここで、インデックス $2k+1, \, 2k+2$ については

インデックス $2k + 1$ : 範囲 $\displaystyle \left[ l, \, \frac{l+r}{2} \right)$ の結果を表す
インデックス $2k + 2$ : 範囲 $\displaystyle \left[ \frac{l +r}{2}, \, r \right)$ の結果を表す

ということになります。従って、更にこの2つのノードについて調べていくことで、目的関数の値を返せるようになります。

まとめ

セグメント木のインデックス $0$ は、範囲 $[0, \, n )$ での値を表しています。このことから、インデックス $k = 0$ から順に範囲 $[ a, \, b )$ との関係を調べていき、以上に示した探索を繰り返していくことで、範囲 $[ a, \, b )$ での目的関数の値を求めることができます。

ソースコード例

たとえば、範囲 $[ a, \, b )$ での最小値を求めるための関数 get_min() は以下のように実装できます。

struct SegmentTree {
public:
  int n;
  vector<ll> node;

private:
  const ll INF = 1e18;

public:
  ll get_min(int a, int b, int k = 0, int l = 0, int r = -1) {
    if (r == -1) {
      r = n; // r = -1 のときは初期化されているので、 r = n として再開する
    }
    if (b <= l || r <= a) {
      return INF; // (i) の場合は最大値を返す
    }

    if (a <= l && r <= b) {
      return node[k]; // (ii) の場合は node[k] の値を返す
    }

    return min(get_min(a, b, 2 * k + 1, l, (l + r) / 2),
               get_min(a, b, 2 * k + 2, (l + r) / 2, r)); // それ以外のときはノード 2k+1, 2k+2 での探索結果のうちの最小値を返す
  }
};

値の更新

元となる数列 ${a}_{0}, {a}_{1}, \cdots, {a}_{n - 1}$ に対して、値の更新があった場合はセグメント木の中身も更新しなければなりません。ここでは、 ${a}_{i}$ の値に変更があった場合について考えます。

まず、セグメント木の中で範囲 $[ i, \, i + 1)$ を表すインデックスは $i + n - 1$ になります。この $i + n - 1$ からセグメント木の上のノードにあたるノードを順に更新していき、最終的にインデックス $0$ に辿り着くまで行っていくことで、セグメント木全体の値の更新も行えます。

先程の図から、インデックス $k$ の下にあたるノードは $2k+1, \, 2k+2$ であるので、

$x = 2 k + 1$ とするとき: $k = \displaystyle \left\lfloor \frac{x - 1}{2} \right\rfloor$
$x = 2 k + 2$ とするとき: $k = \displaystyle \frac{x}{2} - 1 = \left\lfloor \frac{x - 1}{2} \right\rfloor$

ということが言えます。

従って、最初のインデックスを $k = i + n - 1$ とし、 $k = 0$ になるまで $k \leftarrow \displaystyle \left\lfloor \frac{x - 1}{2} \right\rfloor$ として更新していき、その過程でのすべての $k$ について、セグメント木を更新していくことで、全体の更新を行えます。

イメージとしては、セグメント木の一番下から順番に、色を塗られた部分について更新していくことになります。

ソースコード例

セグメント木の各ノードが区間での最小値を扱っているときに、値を更新する update() 関数は以下のように実装できます。

struct SegmentTree {
public:
  int n;
  vector<ll> node;

public:
  void update(int idx, ll val) { // idx が数列 a での位置を、 val が更新後の値を表す
    idx += n - 1; // idx = idx + n - 1 とする
    node[idx] = val; // node[idx] の値を val にする

    while (idx > 0) {
      idx = (idx - 1) / 2; // idx の上のインデックスにする
      node[idx] = min(node[2 * idx + 1], node[2 * idx + 2]); // 下の2つのノードのうちの最小値を node[idx] として更新する
    }
    return;
  }
};

使用例

ここでは、セグメント木の使用例として、「最小値」「最大値」「部分和」のそれぞれの場合についての、「初期化」「取得」「更新」の実装例を示します。

使用例1: 最小値の取得

最小値については、ノード $k$ が範囲 $[ l, \, r)$ を表しているとき、 ${a}_{l}, \cdots , {a}_{r - 1}$ の中の最小値を表していると考えて実装します。

struct SegmentTree {
public:
  int n; // セグメント木を用いて管理する数列の要素数
  vector<ll> node; // セグメント木の頂点となる vector

private:
  const ll INF = 1e18;

public:
  // 初期化
  void init() {
    int sz = 1; // 管理する数列の要素数（2 ^ (m - 1) の形に最終的になる）
    while (sz < n) {
      sz *= 2; // n <= sz となるまで sz の値を2倍する
    }
    n = sz; // n の値を更新する
    node.resize(2 * n - 1); // node の要素数を 2 * n - 1 にする
    fill(node.begin(), node.end(), INF); // node のすべての要素を INF の値にする
    return;
  }

  // 値の更新
  void update(int idx, ll val) { // idx が数列 a での位置を、 val が更新後の値を表す
    idx += n - 1; // idx = idx + n - 1 とする
    node[idx] = val; // node[idx] の値を val にする

    while (idx > 0) {
      idx = (idx - 1) / 2; // idx の上のインデックスにする
      node[idx] = min(node[2 * idx + 1], node[2 * idx + 2]); // 下の2つのノードのうちの最小値を node[idx] として更新する
    }
    return;
  }

  // 値の取得
  ll get_min(int a, int b, int k = 0, int l = 0, int r = -1) {
    if (r == -1) {
      r = n; // r = -1 のときは初期化されているので、 r = n として再開する
    }
    if (b <= l || r <= a) {
      return INF; // (i) の場合は最大値を返す
    }

    if (a <= l && r <= b) {
      return node[k]; // (ii) の場合は node[k] の値を返す
    }

    return min(get_min(a, b, 2 * k + 1, l, (l + r) / 2),
               get_min(a, b, 2 * k + 2, (l + r) / 2, r)); // それ以外のときはノード 2k+1, 2k+2 での探索結果のうちの最小値を返す
  }
};

問題例

使用例2: 最大値の取得

最小値については、ノード $k$ が範囲 $[ l, \, r)$ を表しているとき、 ${a}_{l}, \cdots , {a}_{r - 1}$ の中の最大値を表していると考えて実装します。

struct SegmentTree {
public:
  int n; // セグメント木を用いて管理する数列の要素数
  vector<ll> node; // セグメント木の頂点となる vector

private:
  const ll MIN = -1e18;

public:
  // 初期化
  void init() {
    int sz = 1; // 管理する数列の要素数（2 ^ (m - 1) の形に最終的になる）
    while (sz < n) {
      sz *= 2; // n <= sz となるまで sz の値を2倍する
    }
    n = sz; // n の値を更新する
    node.resize(2 * n - 1); // node の要素数を 2 * n - 1 にする
    fill(node.begin(), node.end(), MIN); // node のすべての要素を MIN の値にする
    return;
  }

  // 値の更新
  void update(int idx, ll val) { // idx が数列 a での位置を、 val が更新後の値を表す
    idx += n - 1; // idx = idx + n - 1 とする
    node[idx] = val; // node[idx] の値を val にする

    while (idx > 0) {
      idx = (idx - 1) / 2; // idx の上のインデックスにする
      node[idx] = max(node[2 * idx + 1], node[2 * idx + 2]); // 下の2つのノードのうちの最大値を node[idx] として更新する
    }
    return;
  }

  // 値の取得
  ll get_max(int a, int b, int k = 0, int l = 0, int r = -1) {
    if (r == -1) {
      r = n; // r = -1 のときは初期化されているので、 r = n として再開する
    }
    if (b <= l || r <= a) {
      return MIN; // (i) の場合は最小値を返す
    }

    if (a <= l && r <= b) {
      return node[k]; // (ii) の場合は node[k] の値を返す
    }

    return max(get_max(a, b, 2 * k + 1, l, (l + r) / 2),
               get_max(a, b, 2 * k + 2, (l + r) / 2, r)); // それ以外のときはノード 2k+1, 2k+2 での探索結果のうちの最大値を返す
  }
};

問題例

E - Smooth Subsequence / AtCoder Beginner Contest 339

使用例3: 区間和の取得

区間和については、ノード $k$ が範囲 $[ l, \, r)$ を表しているとき、 ${a}_{l} + \cdots + {a}_{r - 1}$ の値を表していると考えて実装します。

struct SegmentTree {
public:
  int n; // セグメント木を用いて管理する数列の要素数
  vector<ll> node; // セグメント木の頂点となる vector

public:
  // 初期化
  void init() {
    int sz = 1; // 管理する数列の要素数（2 ^ (m - 1) の形に最終的になる）
    while (sz < n) {
      sz *= 2; // n <= sz となるまで sz の値を2倍する
    }
    n = sz; // n の値を更新する
    node.resize(2 * n - 1); // node の要素数を 2 * n - 1 にする
    return;
  }

  // 値の更新
  void update(int idx, ll val) { // idx が数列 a での位置を、 val が更新後の値を表す
    idx += n - 1; // idx = idx + n - 1 とする
    node[idx] = val; // node[idx] の値を val にする

    while (idx > 0) {
      idx = (idx - 1) / 2; // idx の上のインデックスにする
      node[idx] = node[2 * idx + 1] + node[2 * idx + 2]); // 下の2つのノードの値の和を node[idx] として更新する
    }
    return;
  }

  // 値の取得
  ll get_sum(int a, int b, int k = 0, int l = 0, int r = -1) {
    if (r == -1) {
      r = n; // r = -1 のときは初期化されているので、 r = n として再開する
    }
    if (b <= l || r <= a) {
      return 0; // (i) の場合は関係ないので 0 を返す
    }

    if (a <= l && r <= b) {
      return node[k]; // (ii) の場合は node[k] の値を返す
    }

    return get_sum(a, b, 2 * k + 1, l, (l + r) / 2) + 
               get_sum(a, b, 2 * k + 2, (l + r) / 2, r)); // それ以外のときはノード 2k+1, 2k+2 での探索結果の和を返す
  }
};

問題例

F - Greetings / Codeforces Round 918 (Div. 4)
- 解説記事

あとがき

セグメント木を使う問題と出会ったときは0から実装しているので、個人的なまとめとしてこの記事を書きました。問題例のところについては、使うような問題に出会い、記事を書くなどしたら更新しようかなと思います。

参考にさせていただいた記事

*1:https://kmyk.github.io/blog/blog/2020/03/04/segment-tree-is-not-complete-binary-tree/

yutasの競技プログラミング勉強帖

競技プログラミングの問題についての解説記事を主に書いています。

セグメント木 (Segment tree) 個人的まとめ

この記事の概要

セグメント木とは

セグメント木の基本構造

セグメント木の初期実装

ソースコード例

値の取得

ソースコード例

値の更新

ソースコード例

使用例

使用例1: 最小値の取得

問題例

使用例2: 最大値の取得

問題例

使用例3: 区間和の取得

問題例

あとがき

参考にさせていただいた記事