矮小

井の中の蛙

図解・セグメント木

ARC045#B がセグメント木で殴るタイプの問題だった(頭が良い人は 2 回 imos ると簡単に解ける)んですが僕は頭が悪いのでそろそろ自分でセグメント木を実装できるようになろうと思い立ったのでした.が,調べてみたところいろんなサイトの説明を読んでもなかなか理解できず…….複数のページを見比べて数時間かけてようやく実装できたのですが,どうも文章でしか解説されてないせいで全然理解できなかった気がするのでここではちゃんと図(のようなもの)で示しながら説明していきたいと思います.

セグメント木とは

「配列に対する操作の結果をセグメントごとに保持したもの」を木にしたもの.例えば [2, 3, 8, 4, 1, 5, 7, 6] という配列があったとして,「l 番目から r 番目の要素の中で最小の値は何?」というクエリが何回か投げられるとする (Range Minimum Query: RMQ).これぐらいのサイズの配列なら単純に全部見比べればいいけど,例えば数が 10 億個並んだ配列の中で「1 番目から 10 億番目までの中で……」みたいなことを 10 万回言われると絶対 TLE で落ちる.で,「あらかじめ「全体の中で最小」「右半分で最小」「左半分で最小」……みたいに再帰的に 2 分割しながら結果を持っとけばいいよね」というアイデアなのがセグメント木.上の配列を使うと,

[           1          ]               0
[     2    ][     1    ]       1,             2
[ 2  ][ 4  ][ 1  ][ 6  ]    3,    4,      5,      6
[2, 3, 8, 4, 1, 5, 7, 6]  7, 8, 9, 10, 11, 12, 13, 14 (k)
                          0, 1, 2,  3,  4,  5,  6,  7 (i) (n = 8)

これは完全 2 分木になっていて,根ノード=min(左の子,右の子) になっていることが分かる.さらに完全 2 分木にすると何が嬉しいかというと,これを 1 次元の配列で扱えるようになる.上図右は左の 2 分木をルートから BFS で 0-indexed の添字を付けたもので,左の子=2 * 親 + 1,右の子=2 * 親 + 2になっている.また,元の配列の長さが 8 で,元の配列の 0 番→ 2 分木配列の 7 番,1 番→ 8番,……という関係なので,元の配列の長さ (8) を n とすると,元の配列の i 番目の要素は 2 分木配列では k = i + n - 1 番目になっている.これで 2 分木配列と元の配列の間でインデックスの変換ができるようになった.

値の更新

セグメント木で出来るのは配列内の 1 要素の更新と,木のある範囲に対する reduce 操作である(max, min, gcd, sum など).要は親=操作(左の子,右の子) が常に成り立てばいい.完全 2 分木では,k 番目のノードの親ノードの番号は (k - 1) / 2 (切り捨て)で求まる.ということで,値の更新の擬似コードは以下のようになる:

def update(元の配列の番号, 新しい値)
  2分木内の番号 k = 元の配列の番号 + n - 1
  2分木配列[k] = 新しい値

  while k > 0:
    # 1個上の親に上がる
    k = (k - 1) / 2
    # 子を使って親の値を更新する
    2分木配列[k] = 操作(2分木配列[2 * k + 1], 2分木配列[2 * k + 2])

値の取得

[           1          ]               0
[     2    ][     1    ]       1,             2
[ 2  ][ 4  ][ 1  ][ 6  ]    3,    4,      5,      6
[2, 3, 8, 4, 1, 5, 7, 6]  7, 8, 9, 10, 11, 12, 13, 14 (k)
                          0, 1, 2,  3,  4,  5,  6,  7 (i) (n = 8)

例えば,上図で「元の配列で言う 2 番から 7 番までの中で最小の値が欲しい」と言われたら,2 分木配列の添字を見て,直感的に「2 分木配列で言う 2 番,4 番のうち小さい方を返せば良さそう」と思うだろう.思うことにしよう.で,これはどういう直感に基づいているかというと「2 分木配列のノードの中で,欲しい範囲に完全に包含されておりかつもっとも大きい範囲を担当するノード」を返せばいいことになる(これはたぶん iwi さんのスライド 53 枚目が一番分かりやすい).そのためには,2 分木配列の k 番目のノードが元の配列で言う何番から何番までを担当しているかを教えながら再帰的に 2 分探索する必要がある.ということで擬似コードは以下の通り:

# この補助関数を getValue(a, b) = __getValue(a, b, 0, 0, n - 1) と呼べば
# 上から順に探索できる
def __getValue(元の配列の左端 a, 同右端 b,
               2分木配列中の番号 k, k番が担当している元の配列の左端 l, 同右端 r):
  if k番が担当している元の配列の範囲 l:r が a:b に完全に含まれている:
    return 2分木配列[k]

  if l:r が a:b から完全に外れている:
    return 単位元(minならint.MAX,sumなら0とか)

  return 操作(__getValue(a, b, 2 * k + 1, l, (l + r) / 2),
             __getValue(a, b, 2 * k + 2, (l + r) / 2 + 1, r))

この関数では,特に l と r が 2 分木内の添字ではなく元の配列の添字を指していることに注意.また,他の例ではよく半開区間で考えると左の子の右端と右の子の左端が同じになって分かりやすいみたいなことが書いてあったけど,僕にはそれをプログラム内でどう扱えばいいのか分からなかったので普通に閉区間にした.ぶっちゃけ開区間とか閉区間の概念を理解していればどっちでもいいと思う.

プログラム例を挙げるだけでは再帰の処理が何をやっているのか分かりづらいので,例を出して処理を追ってみよう.例えば,getValue(2, 7) とすると補助関数の中では変数が以下のように動いている.

             0
     1,             2
  3,    4,      5,      6
7, 8, 9, 10, 11, 12, 13, 14 (k)
0, 1, 2,  3,  4,  5,  6,  7 (i) (n = 8)
2, 3, 8,  4,  1,  5,  7,  6 (元の配列)
      a                   b
l                         r  k = 0     => min(4, 1) => 1
l         r                  k = 1   => min(MAX, 4) => 4
l  r                         k = 3 => MAX
      l   r                  k = 4 => 4
              l           r  k = 2   => 1

2 分木配列の初期化

2 分木配列を初期化するときは,元の配列の長さ以上で最も小さい 2 の冪を求めたあと(完全 2 分木は末端のノード=葉の数が必ず 2 の冪になる),それ * 2 - 1 の長さで中身を単位元で埋めた配列を用意して,k = i + n - 1 番目から順に update していけばいい.たぶんコンストラクタに元の配列を渡してコンストラクタの中で全部やればいいと思う.

まとめ

  • セグメント木は,大きい配列の中の部分配列に対する reduce 操作を完全 2 分木を使って出来るようにしたもの
  • 値を更新するときは 2 分木を下から上に辿り,値を求めるときは 2 分木を上から分割しながら辿っていけばいい

間違ったこと・分かりづらいことを言ってるところがあったらコメントください.随時修正します.

ちなみに

これを Python 3 で書いて提出したら TLE になった.悲しい.

Submission #525117 - AtCoder Regular Contest 045 | AtCoder

参考

www.slideshare.net

d.hatena.ne.jp

d.hatena.ne.jp

taku-k.hatenablog.com