精進してますか?
(セコムしてますか?と同じ口調で)
こんにちは。六月と申します。
先月に公開させて頂いた「プログラミング初心者の40代おじさんが1年かけてAtCoder緑になった話(色変記事)」は多くの方にお読み頂き、ありがとうございました。これを機に競プロ(競技プログラミング)を始めてみたいといった声も頂き、大変うれしく思っております。
とはいえ、本来Qiitaは技術的な知見をシェアする場だと心得ています。色変記事は言ってみれば私の身の上話ですので、それだけにとどまらず、私も何か技術的な記事を書いてみたいなと思っていました。
それを受けて、最近知見を得た「セグメント木」(通称セグ木。以下、本文中ではセグ木と表記します)というデータ構造について書いてみたいと思います。
セグ木は、競プロ上達にあたり避けて通れないデータ構造らしく、私もヒマを見つけて勉強しておりました。
しかし、やりたいことはなんとなくわかるが、なぜこんな回りくどいやり方でデータ管理しなければならないのか、という疑念がなかなか拭えませんでした。そんなこともありなかなか手が進まず、基本のキである一点更新・区間取得のシンプルなセグ木を実装するだけなのに、勉強し始めから実装完了まで約3ヶ月もかかってしまいました。
でも、セグ木の動き方とその計算量削減効果をきちんと理解できたとき「私がアホだった…。このやり方を最初に考えた人は天才だわ…」とめちゃくちゃ感動しました。
また、自分で実装してみると、本当に多くの学びがあります。
ということで、今回はこの3ヶ月間の試行錯誤でセグ木を枯らせまくってきた私が、これからセグ木を勉強する人のために、セグ木の理解→実装開始までのつまずきポイントに落ちている石を、可能な限り取り除きます。
そして、これからセグ木を勉強する人には、私が3ヶ月かかった旅路を3日で走り抜けてもらいたいです。
それがこの記事を書いた目的です。
そして、これを読み終わった後は、ぜひ自分だけのセグ木を育ててみてください。自分がゼロから書いたセグ木が動くと結構感動しますよ~。
これが俺たちのガーデニングだ!
それでは行ってみましょう。
(本文はなるべく普遍的な記述を心がけておりますが、私がPythonしか書けないため、サンプルのソースコードがPythonです。他言語をご利用の方、その点はご容赦ください)
なぜセグメント木が必要なのか
セグ木とは「区間に対する処理を、高速に行うデータ構造」です。これだけでは何のことかよくわかりませんね。
そこで、サッカーのJリーグで、ゴールを決めた人のデータが試合後に毎回与えられる状況を考えてみましょう。
たとえば「2021年の得点王はゴールを何点決めましたか?」と聞かれたら答えるのは簡単です。もう全ての試合が終わっているからです。全得点を記録してソートし、一番ゴール数が多かった人の得点を答えるだけでOKです。
ところが、次のような質問が来ると答えるのが少し難しくなってきます。
-
毎週末の試合が終わったら、その時点の得点王が何点決めているか知りたい
-
毎週末の試合が終わったら、その時点の得点王が何点決めているか知りたい。ただし関東に拠点を置くチームだけで集計してほしい
-
あるチームだけで、チーム内の得点王が何点決めているか、毎試合終了ごとに知りたい
なぜ難しいかと言うと、まだリーグ全試合が終わっていない状況だと、これらの質問に答えるには、試合終了ごとに集計結果をまとめて、毎回ソートしなければならないからです。
現実のJリーグは20チーム(J1)だけなので、力づくで集計できるかも知れませんが、もしJリーグの所属チームが膨大にあり、管理対象となる選手が10万人もいたとしたら?
問い合わせの都度、10万ものデータを毎回ソートしていたら、さすがにコンピュータでもそれなりに時間がかかります。
さらに競プロのように、どんどん値が変わる10万の管理対象データに対し次々と問い合わせが飛んできて、全質問に対する正しい回答を2秒以内にくれ、みたいな状況だったら?
このような過酷な要求に耐えるのがセグ木というデータ構造です。
大ざっぱに言うと、管理対象のデータがまだ変化している途中なのに、その最中で何回も「今どんな感じ?」と聞かれてしまうような状況にあった場合、セグ木の出番となります。
セグメント木はどうやって作る?
まず結論から言うと、セグ木は管理対象となるデータ件数の2倍のサイズの配列で、効率的にデータを管理します。
ちなみにその中で実際に使うのは、データ件数の2倍から1を引いた部分です。つまり配列のインデックス1個分は使いません。
データ件数が8件の状況を考えてみます。この場合、セグ木のサイズは8*2==16が必要ですので、初期値として0を16個格納した配列を作ります。
[0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0]
そして、わかりやすくするためサッカーを例にとり、8人の選手の得点数を管理し、その時点での得点王が何点決めているか答えるセグ木を作ってみましょう。
イメージしやすくするため、1~8の数字にそれぞれ選手を割り当てました。
(選手名は架空のものです!)
選手番号と選手名の対応
1 カズ
2 オオサコ
3 ゴン
4 ザキオカ
5 カマモト
6 タケダ
7 サワ
8 ラモス
とします。
そして、ここがセグ木の大事な部分になりますが、管理対象のデータ数が8、つまり配列のサイズが16となるセグ木では管理データの1番目と配列のインデックス8が対応します。そして管理データの2番目は配列のインデックス9が対応します。
わかりにくいのでまとめます。下記のようになります。
- 選手番号とセグ木配列のインデックスとの対応関係
選手番号 | インデックス |
---|---|
1(カズ) | 8 |
2(オオサコ) | 9 |
3(ゴン) | 10 |
4(ザキオカ) | 11 |
5(カマモト) | 12 |
6(タケダ) | 13 |
7(サワ) | 14 |
8(ラモス) | 15 |
つまり、具体的には、選手番号4のザキオカ選手が得点を決めたら、配列のインデックス11にまずその情報を反映させることになります。
7のサワ選手が得点を決めたら、配列のインデックス14にその情報を反映させます。
つまり、配列を半分で割った後半に各選手の得点データが入ります。
配列の前半は、後半部分の得点データを効率的に管理するために使われます。
これが、セグ木の構築時にデータ件数の2倍のサイズの配列を作る理由です。
なお、前述の通りセグ木で実際に使うのは、正確には配列のうちデータ件数の2倍から1を引いた部分です。
今回のセグ木の場合、配列の0番目は使わず、1~15番目を使います。そのため、配列0番目の値は仮で入れているだけです。今回は便宜的に0を入れていますが、ダミーだということが自分でわかるならば、違う値を入れても大丈夫です。
また、配列の0番目から使い始めることも可能です。その場合は、0番目~14番目を使い、末尾の15番目が不要となります。
なお、配列を0番目から使うことを0-indexedと呼び、1番目から使うことを1-indexedと呼ぶそうです。どちらにするかは人それぞれの流儀によりますが、熟練したプログラマーの方ほど0-indexedを好む印象を受けます。あと、他の人のソースコードを見ていても、0-indexedをサラッと使いこなしている人はカッコよく見えますね。でも私は1-indexedが好きです。だって0番目とか言われてもよくわかんないしさ…。
ところで、私が今回の記事で1-indexedを採用した理由は他にもあります。後述します。
得点の情報をセグメント木に反映させるには?
次に、選手が得点を決めたらセグ木に反映させることを考えてみましょう。
n番の選手がg点決めたら
n g
の形式で入力が標準入力から与えられるということにします。
さっそく得点の情報が飛んできました。
3 2
3番のゴン選手が2ゴール決めたようです!
ゴンゴール!
(選手名は架空のものです)
このとき、セグ木は次のように更新されていきます。
まず初期状態はこちらです。得点者はまだ誰もいませんね。
3番の選手が2得点なので、まず最初に選手番号3に対応するインデックス10に2を書き込みます。
次に、その真上のインデックス5に2を書き込みます。
次に、その真上のインデックス2に2を書き込みます。
最後に、一番上のインデックス1に2を書き込みます。
最終的に、配列では4箇所が更新されました。
また得点の情報が標準入力から与えられました。
7 3
7番のサワ選手が3ゴール決めたようです。
ハットトリック!
さすがバロンドール…。
(繰り返しますが選手名は架空のものです)
するとセグ木は次のように更新されていきます。
最初に、選手番号7に対応するインデックス14に3を書き込みます。
次に、その真上のインデックス7に3を書き込みます。
次に、その真上のインデックス3に3を書き込みます。
最後にインデックス1です。ここには先程の3番(ゴン選手)の2得点の情報が既に書き込まれていますが、7番(サワ選手)の3得点の方が数が大きいので、インデックス1の値を3に更新します。
なんとなく仕組みが見えてきましたか?
上下のインデックスを求める方法
セグ木では、頂点のインデックスは真下にある2つのインデックスの値を管理しています。
そのインデックスはそのまた真下にある2つのインデックスの値を管理しています。
これを繰り返すと、最終的に選手の得点を直接管理している最も下のインデックスにたどり着きます。
そして、選手の得点が更新されると、真上にあるインデックスと値を比較し、大きければ値を更新して次々と上へ上げていきます。そうすることで、全選手の中で一番得点の多い人の値がインデックス1に反映されるようになります。
これがセグ木が値を更新する仕組みです。
なお、あるインデックスを2で切り捨て除算すると、真上のインデックスがわかります。これもセグ木の大きな特徴です。
先程の選手番号7番(サワ選手)の値がどのように更新されたかもう一度見てみましょう。
まず選手番号7番に対応するインデックス14に3得点の情報を書き込みました。
次に、14//2==7で、真上のインデックスが7だとわかります。そしてインデックス7の値を更新します。
次に、7//2==3(余りの1は切り捨て)で、真上のインデックスが3だとわかります。インデックス3の値を更新します。
最後に、3//2==1(余りの1は切り捨て)で、真上のインデックスが1だとわかります。インデックス1の値を更新します。
逆に下に下がるときは、インデックスを2倍すれば左下、2倍+1すれば右下のインデックスに対応することがわかります。
さらに言えば、2以上のインデックスについては、偶数であれば左側、奇数であれば右側になります。
先ほど、このセグ木では「1-indexedを採用した」と述べましたが、それはこの真上または真下のインデックスを求める計算を簡単にしたかったためです。
この仕組みを活用することで、管理するデータに変動が生じても、わずかな計算回数で膨大な区間の最大値を管理できます。
なお、このセグ木では、ある選手の得点の情報が与えられた際、whileループを3回まわすとインデックス1にたどり着きます。
なんだよ、たった8人の選手の得点を管理するだけなのに得点情報1回あたりループを3回もまわすのか、どこが効率的なんだと思うかも知れません。
しかし、セグ木の備える効率性は、選手数が大きくなればなるほど効いてきます。
冒頭で、選手が10万人もいたらどうやって管理するのか、と述べましたが、セグ木を使えば10万人の選手がいても、わずか17回ループをまわすだけで最下のインデックスから頂点のインデックス1にたどり着けます。
たとえばちょうど10万人の選手がいたとしましょう。
すると、セグ木の配列では、一番数字の大きい選手番号100000の選手の得点を直接管理するインデックスは231071になります。
先ほど説明したように、あるインデックスの真上のインデックスは2で切り捨て除算すれば求められます。実際にやってみましょう。
1回目 231071//2 == 115535
2回目 115535//2 == 57767
3回目 57767//2 == 28883
4回目 28883//2 == 14441
5回目 14441//2 == 7220
6回目 7220//2 == 3610
7回目 3610//2 == 1805
8回目 1805//2 == 902
9回目 902//2 == 451
10回目 451//2 == 225
11回目 225//2 == 112
12回目 112//2 == 56
13回目 56//2 == 28
14回目 28//2 == 14
15回目 14//2 == 7
16回目 7//2 == 3
17回目 3//2 == 1
ということで、17回で1にたどり着きます。
このように、セグ木を使ってデータ管理すれば、いつどのタイミングで「今の全区間最大値を教えてくれ」と言われても大丈夫です。既に見てきたように、セグ木のデータ更新方法であれば、インデックス1に常に全区間最大値が反映されているからです。
つまり、もし10万人の選手がいたとしても、10万ものデータが格納された配列をいちいちソートして調べる必要はなく、セグ木のインデックス1を見れば現在の得点王が何点決めているのかすぐにわかります。
これがセグ木の効果です。
任意の区間の最大値を教えてくれと言われたら?
セグ木が管理データを更新する仕組みがわかったところで、次は任意の区間の最大値を求めることを考えてみます。
たとえば、何試合か進み、現在の得点数が以下のような状況だったとしましょう。
現在、選手番号1のカズ選手が得点ランキング1位です。
さすが俺たちのキングカズ!
そして選手番号8のラモス選手が2点差で追っています。得点王はこの2人のどちらかで間違いなさそうです。
(選手名は架空の…(以下略))
そこで残る選手、具体的には選手番号2~7の選手の中で、いま一番得点の多い人が何点決めているか知りたいとします。
セグ木はこのような、任意の区間に対する問い合わせにも対応します。いま例として挙げている2~7だけではなく、2~5でも、3~7でも、区間が連続していればどの区間でもOKです。
さて、そういった任意の区間に対する問い合わせがあった場合、セグ木はどのように動くか確認しましょう。
まず、区間に対する問い合わせがあった場合、トーナメント戦のように問い合わせ区間の左端と右端がぶつかるまで上に上げていくことを考えます。
今回は選手番号2~7の選手で検討します。
区間左端の選手番号2はインデックス9となります。
反対に右端は選手番号7はインデックス14です。
ここから説明する内容もセグ木の大きな特徴です。
区間に対する問い合わせがあった場合、左端で今見ているインデックスが奇数だった場合、不戦勝のような形で、無条件で上に勝ち上がり、2回戦に進めます。
選手番号2つまりインデックス9は、奇数なのでそのまま2回戦のインデックス4の位置に進めます。実際、隣のインデックス8(黒く塗りつぶしたところ)は選手番号1に対応するため、区間問い合わせの範囲外であることがわかります。
右端の場合は左端とは真逆の考え方となり、今見ているインデックスが偶数だった場合、不戦勝のような形で無条件で上に上げることができます。実際に、選手番号7つまりインデックス14の隣のインデックス15は選手番号8に対応するため、やはり区間問い合わせの範囲外です。
ここはめちゃくちゃ大事なので、説明が悪かった場合は申し訳ないのですが、よく理解できなかった場合は区間に対する問い合わせの左端と右端の数字を3と6とか、2と5とかいろいろ入れ替えて上に上るときの動き方をイメージしてみてください!
どの位置を切り取っても、左端のインデックスが奇数だった場合、左隣の偶数のインデックスは区間問い合わせの範囲外になることがわかります。
また、この性質は上の段に登っても変わりません。
逆に、右端のインデックスが偶数だった場合は、右隣に来る奇数のインデックスが区間問い合わせの範囲外となります。
セグ木のこの性質を利用することで、うまいこと問い合わせ対象外の範囲を避けながら、トーナメント戦を勝ち上がるがごとく、ループを1回まわすごとに、左端と右端をそれぞれ1つずつ上に上げていくことができます。
言ってること伝わってるかな?
頼む伝わってくれー!
次の動きを見てみます。
左端の勝者である選手番号2の18点はいま、インデックス4の位置にいます。今度はインデックスが4で偶数なので、不戦勝とならず、隣のインデックス5と戦う必要があります。
隣のインデックス5は19点なので負けてしまいました。そしてインデックス5の19点が勝者として上のインデックス2の位置に上がり、新しい左端として振る舞うことになります。
右端は今インデックス7の位置にいます。今度はインデックスが奇数なので、やはりこちらも不戦勝とならず、隣のインデックス6と戦います。
隣のインデックス6は13点なので、インデックス7が勝ちました。こちらはインデックス7の17点が変わらずそのまま右端の勝者として、上のインデックス3の位置に上がります。
次の動きを見てみます。
左端は今インデックス2の位置にいます。偶数なので不戦勝とならず、隣のインデックス3と戦う必要があります。
そして右端はいまインデックス3の位置にいます。奇数なので不戦勝とならず、隣のインデックス2と戦う必要があります。
ここで、左端の勝者と右端の勝者の直接対決となりましたので、これが最後の戦いです。そして19点>17点ですので、インデックス2の位置にいる左端の勝者が最終的な勝者となりました。
そしてこの19点という数字が、先ほどの「選手番号2~7で一番多い人の得点は何点か」という問い合わせに対する答えとなっていることがわかります。
これが、セグ木が区間に対する問い合わせに対応する仕組みです。先ほど見たように、データが10万あっても、17回のループでインデックス1までたどり着けますので、わずかな計算量で左端と右端を直接ぶつけることができます。
そうすることで、わざわざ大きなサイズの配列をソートしなくても、区間の最大値がわかるのです。
なお、この区間取得のやり方は他にもあります。
たとえば、先ほどの選手番号2~7ということは、インデックス9、インデックス5、インデックス6、インデックス14の中から最大値をとることと同じです。
このように、問い合わせ内容からその計算に必要なインデックスを直接割り出し、その値を比較して答えを求める方法もあり、そちらのほうが高速に動作します。
今回私が採用しているトーナメント方式は実装が簡単である一方で、やや動作が遅めです。
(2024年12月13日追記)
上述した「その計算に必要なインデックスを直接割り出す」という処理を実装させる問題が、AtCoder Beginner Contest 349のD問題で出題されました。セグメント木の区間取得では、当記事で紹介したトーナメント方式のやり方よりも、この問題が要求する実装方法の方が一般的です。下にリンクを載せますので、解説含め一度目を通しておくことをおすすめします。
AtCoder Beginner Contest 349
D - Divide Interval
データ数が中途半端だった場合、セグメント木の配列の要素数はどうする?
セグ木では、あるインデックスは真下にある2つのインデックスを管理します。そのため、管理するデータ数が、2、4、8、16といったように、2を倍々していった数字だった場合、配列の後半にきれいに収まります。
では、データ数が5とか、中途半端な数字だった場合はどうしたらいいのでしょうか。その場合は、ダミーのデータを入れておけば大丈夫です。
実際に見てみましょう。先ほどの得点数を管理する例では8人の選手がいましたが、選手が5人だった場合を考えてみます。するとセグ木はこうなります。
選手が5人しかいないので、元々選手番号6~8の得点数を管理していた箇所には、計算結果に影響を与えないよう、ダミーとして0を入れています。これでもセグ木としては問題なく動きます。
ただこれだと、データ数が中途半端だった場合、初期化時にセグ木のサイズは単純にデータ数の2倍、というわけにはいかず困ってしまいます。初期化する際の配列の要素数はどのように判定したらいいのでしょうか。
ここは、ソースコードを見てもらった方がわかりやすいかも知れません。今回のセグ木で、データ数に応じてセグ木を初期化する部分のPythonのソースコードはこちらです。
#監視対象のデータ数Nに合わせてセグメント木を初期化。
#変数tmpの値を1→2→4→8、と倍々にしていき、
#Nを超えたところでストップ。すると、その値を2倍した
#数字がセグメント木に必要な配列の要素数になる。
tmp = 1
INF = 0 #←初期化したときに入れる値は問題によって変えましょう
while tmp<N:
tmp*=2
size = tmp*2
data = [INF]*size
つまり、
データ数が2だったら、セグ木の配列の要素数は4
データ数が3~4だったら、セグ木の配列の要素数は8
データ数が5~8だったら、セグ木の配列の要素数は16
データ数が9~16だったら、セグ木の配列の要素数は32
ということになります。
ただ、この方法は実装が楽というメリットがある一方で、管理するデータ数によってはセグ木の実体である配列内に、実際には使わないムダな部分がたくさん存在してしまうというデメリットがあります。
データ更新方法や区間処理のやり方を工夫することで、ムダな領域をなくすようなセグ木の組み方もできるようですので、次のステップとしてぜひチャレンジしてみてください。
今回のセグメント木のソースコード
まとめとして、今回のセグメント木のソースコードを載せておきます。どうも私はシンプルな記述が苦手で、同じような処理でも冗長になってしまうクセがあります。その点はご容赦ください。
また、本文の内容を要約して該当する処理の部分にコメントとして書きました。処理の過程がよくわからなかった場合は本文に戻って見比べてみてください。
N = int(input()) #選手の数
#=====RmQ/一点更新・区間最大値を求めるセグメント木=====
#監視対象のデータ数Nに合わせてセグメント木を初期化。
#変数tmpの値を1→2→4→8、と倍々にしていき、
#Nを超えたところでストップ。すると、その値を2倍した
#数字がセグメント木に必要な配列の要素数になる。
tmp = 1
INF = 0 #←初期化時に入れる値は問題に合わせて変えましょう
while tmp<N:
tmp*=2
size = tmp*2
data = [INF]*size
#得点時に選手番号kの得点数を加算する。
def updata(k,a):
kk = (k-1)+tmp
data[kk] += a
#インデックス1に届くまで隣のインデックスと値を
#比較し大きい方をどんどん上に上げる
while kk>1:
if kk%2 == 0:
data[kk//2] = max(data[kk],data[kk+1])
else:
data[kk//2] = max(data[kk],data[kk-1])
kk = kk//2
return
#選手番号l~rの最大得点数を計算
def query(l,r):
left = (l-1)+tmp
right = (r-1)+tmp
if left == right:
return data[left]
#区間に対する処理をトーナメント戦に見立てて
#左端と右端を上にどんどん上げていく
winner_left = data[left] #←左端の暫定勝者
winner_right = data[right] #←右端の暫定勝者
while left//2 != right//2:
#今見ている左端のインデックスが偶数だった場合のみ
#右隣のインデックスの値と戦う
if left%2 == 0:
if left+1<=right:
winner_left = max(winner_left,data[left+1])
left = left//2
#今見ている右端のインデックスが奇数だった場合のみ
#左隣のインデックスの値と戦う
if right%2 == 1:
if right-1>=left:
winner_right = max(winner_right,data[right-1])
right = right//2
#左端の暫定勝者と右端の暫定勝者が直接対決になったら
#そのmaxを取り、値を返す
return max(winner_left,winner_right)
#=====セグメント木ここまで=====
updata(3,2) #選手番号3が2得点
updata(7,3) #選手番号7が3得点
print(query(2,7)) #選手番号2~7の最大得点が何点か出力
print(data) #セグメント木の実体である配列を出力
私にはそのスキルがなくまだできませんが、こういうのをクラスにまとめたりできる人はすごくカッコいいですよね。
私も次はそれにチャレンジしたいと思っています。
おまけ・作ったセグ木の動作を確認できる問題集
今回のセグ木は、一点更新・区間最大値を求めるものですが、maxのところをminに変えれば、区間最小値を求めるセグ木になります。このように、自分で実装してみると問題に合わせて簡単にカスタマイズできるのがメリットです。なにせ自分で書いたものですからね。
まだセグ木をいじった経験のない人はぜひチャレンジしてみてください。
また、セグ木には今回のように1点更新ではなく、区間の値をまとめて高速に更新できる遅延評価セグメント木というスペシャル版があります。
じゃあなんでそれをこの記事でやらないのかって?
まだ書けないんだよ言わせんな恥ずかしい…
最後におまけとして、作ったセグ木がきちんと動くか確認できる問題をいくつか紹介して終わりにしましょう。
Range Minimum Query (RMQ)
AOJの問題です。一点更新・区間最小値を求めるセグ木を作ります。
Range Sum Query (RSQ)
同じくAOJの問題です。一点更新・区間和を求めるセグ木を作ります。
B - Fenwick Tree
AtCoder Library Practice Contestの問題です。本来はFenwick木(BIT木)の練習のための問題ですが、セグ木でも解けます。同コンテスト内にはセグメント木の練習のための問題もありますが、そちらは難易度が高いため、まずはこのFenwick木用の問題で動作確認してみることをオススメします。
まとめ
長くなりましたが、ここまでお付き合い頂きありがとうございました!
完全に余談ですが、最近、私の近所では庭の木が大きくなり、お隣さんの敷地に枝が伸びていってしまう越境問題が頻発しており、町内会でもよく注意喚起しています。でも私は大丈夫です。だって私が生やしているのはセグ木ですからね…。これがメタバース時代の庭いじりというものであります(違う)。
というわけで、ここでお別れです。
またいつかお目にかかりましょう。