とりになりたいかものはし

プログラミング(主に競プロ)に関するひとりごとなど

セグ木に平衡二分木を載せよう!

 セグメントツリーってすごいですよね!原理はシンプルでありながらも驚くほどの汎用性を持ち、 「区間の最小値」みたいな典型問題はもちろん、DPの漸化式に放り込んでみたり、結合法則が成り立つ構造体を入れてみたり、 果てには集合の簡易的な表現に使ってみたりと、競プロerの生活に欠かせないデータ構造です。 (筆者も先日のPCK予選で2回ほどお世話になりました)

 今回はこのセグメントツリーと同様に有名な、でも標準ライブラリに実装されている機能で普段事足りているから実装したことがない人も多いかも?という 「平衡二分木」をセグメントツリーに載せてみたいと思います!でもその前に、少し予習をしておきましょう

プログラミングコンテストでのデータ構造
プログラミングコンテストでのデータ構造 2 ~平衡二分探索木編~

 2つともiwi先生のスライドで、一つ目は後半にセグメントツリーの説明が、2つ目はまるまるtreapによる平衡二分木の実装が入っています。
上のスライドを読むのもだるい!っていう方は以下の2つのポイントだけ理解していればOKです。 (だるくなってきたので以下からセグメントツリーを「セグ木」と呼びます)

・基本的なセグ木の書き方

さすがにこれはほかの方々がすでに何回も紹介されていると思うので、このブログで取り上げるべき話題ではないと思われます。

・平衡二分木は何ができるか?

 「木」と言っておきながら、平衡二分木は「常にソートされている数列」みたいに考えておくのが一番無難だと思います。 例えば、以下のすべてのクエリをO(log N)で処理できます。(N=木の大きさ)
①値の挿入(insert)・値の削除(erase)
②k番目の値の出力
③値の検索
④ある値未満の数列 と ある値以上の数列 への分割(split) ex) {1,2,4,5,7}のとき、5を境目に分割して {1,2,4} {5,7} にする
⑤ある値未満の数列 と ある値以上の数列 の統合(merge)
⑥数列内のすべての値の和、値の個数、値の最大値最小値etc...
  ⑤の統合(merge)についてですが、あらゆる平衡二分木を統合できるわけではないことに注意です!  例えば、{1,2,4}{5,7}は、どの左の要素よりも右側の要素のほうが大きいので統合可能ですが、{1,2,5}{4,7}などは無理です。(もちろんerase/insertを使い、値を一つずつ引っぺがしてはもう片方に入れていくとO(N log N)で可能です)  一見あまり役に立たなさそうなmergeですが、内部実装などで使う場合があるので機能の一部に入っています。 実はC++の標準ライブラリに実装されている「std::set」および「std::map」は上記のうち①③しか実装されていません! また、g++の拡張機能であっても上の2つに加えて②しか使えないそうです。 なので④⑤⑥(とりわけ⑥)の機能が欲しい今回のブログのテーマにおいては、平衡二分木は自前で全部実装したものを用意しておいてください。

本題1

以下の問題を考えてみましょう。

「横幅がN、縦に無限に広い長方形がある。大きさ1×1のマス目に区切られており、各マス目のコストは最初0である。 今から2種類のクエリを合計Q個処理せよ。
クエリa:(l,r,a,b,x)が与えられ、幅[l,r)、高さ[a,b)の長方形内に存在するすべてのマス目のコストにxを足す。
クエリb:(x,y)が与えられ、座標(x,y)に存在するマス目のコストを出力する。 N,Q≦10^5」

例: N=5の時
クエリa,(1,5,2,4,5)
クエリa,(2,6,3,5,-1)
クエリb,(0,0)→0
クエリb,(1,2)→5
クエリb,(3,3)→4
クエリb,(5,4)→-1

f:id:platypus999:20170920212445j:plain

解法: セグ木のすべてのノードにpair<int,int>のような、二つの値を要素に持つ平衡二分木を載せます。 最初はすべての平衡二分木を空にします。

クエリaが来た時、
①[l,r)区間に対応するノードを通常のセグメントツリーのように列挙します。セグ木の性質により、ノードの数は高々$O(log N)$個であることがわかります。
②上のノードに対し、ひとつずつに(a,x)と(b,-x)、2つの値を挿入します。 一回の挿入にかかる計算量は$O(log Q)$ですので、全体としては$O(log N * log Q)$でクエリは処理できます。

クエリbが来た時、
①すべてのノードのうち、xを含む(区間の左端≦x≦区間の右端 が成立する)区間を列挙します。セグ木の性質により、これの個数も高々O(log N)です。
②すべてのノードの平衡二分木を見ます。それぞれの木に対し、値が(y,∞)以下の要素の和を求め、その第二値(C++で言うとpairのsecond)を足し合わせます。平衡二分木の分割・和の取得の計算量は$O(log Q)$のため、全体としては$O(log N * log Q)$で処理できます。 以上の処理で出てきた値がクエリbの答えになります。

 以上のアルゴリズムを使うと、Q個のクエリに対し、$O(Q * log Q * log N)$の計算量で処理できます。 $O(N)=O(Q)$の場合、$O(N log^2 N)$となるため、愚直にやった場合の$O(N^2)$に比べて改善がされています。

証明: なぜ上のアルゴリズムによって正しい結果が得られるのでしょうか? まず、より簡単な問題から分析していきましょう。

「無限に長い数直線がある。今から2種類のクエリを合計Q回処理せよ。
クエリa:(l,r,x)が与えられ、[l,r)の区間に一様に値xを足す。
クエリb:(k)が与えられ、数直線の場所kの値を出力する。 Q≦10^5」

 これはBITとクエリ先読みによる座標圧縮を使えば簡単に求まる問題です。ただ、メモリを節約するために、またinteractiveな問題にも対応するために、平衡二分木を使うこともできます。
まず、 空の平衡二分木を1個持ちます。

 クエリaが来た時、(l,x)と(r,-x)、2つの値を挿入します。(計算量は$O(log N)$)
クエリbが来た時、平衡二分木中のうち、値が(k,∞)以下のすべての要素の和を求めます。この時、この第二値がクエリbの答えです。(分割と値の総和の取得で処理できるので、これも計算量は$O(log N)$) なぜならば、クエリbのkについて、すべてのクエリaの[l,r)を見たとき、 k<lならば何も足されないので0 l≦k<rの時、(l,x)は足されますが(r,-x)は足されないのでx r≦kの時(l,x)、(r,-x)が両方足されるのでx+(-x)=0 となるからです。

 この性質を応用して、2次元に拡張しましょう。 二次元平面を列ごとに分割し、1列ごとを平衡二分木で管理しましょう。 幅[l,r)高さ[a,b)の長方形に一様に値xを足すときは、 l≦i<rであるすべてのiについて、 「i番目の平衡二分木に(a,x)と(b,-x)を挿入する」 という処理をすれば、各クエリはO(N log Q)で処理できそうです。

 さらに、[l,r)という区間において、l≦x<rを満たす一つ一つのxを見ていくのは無駄です。 横の区間をセグ木のように持ってみましょう。次に、各区画に平衡二分木を持たせます。 こうして、クエリaが来た場合は、 区間[l,r)は[l1,r1),[l2,r2),[l3,r3),[l4,r4),...,[ls,rs)というs個の区間に分けることができます。 セグ木の性質上、sは$O(log N)$であることがわかります。 そして、これらs個の区画のそれぞれの平衡二分木に(l,x),(r,-x)を足します。 計算量は全体で$O(log N * log Q)$です。 さて、クエリbが来たときはどうすればよいでしょうか? これについては、xを含んでいるすべての区間(l≦x<rであるような[l,r))を見ます。 セグ木の性質より、この条件を満たす区間の個数は$O(log N)$です。 そして、各区間の平衡二分木に対し、「(y,∞)以下の値の和」を求め、その総和が答えになります。

実践

データ構造はACしないと意味がない!というわけで2問ほど

NCR CodeSprint(HackerRank) Points and Fences

問題文
 すこし考察をすると、この問題は「長方形区間に値をxorする」「ある点の値を取得する」 の二つのクエリを高速に処理すればよいことがわかります。
これは単にセグ木に平衡二分木を載せて、平衡二分木中のすべての値のxor値を取得する機能を追加すれば簡単に処理できます。 計算量も$O(q * log n * log q)$と十分高速です。 (実はこの問題はセグ木にBITを載せたデータ構造でも解くことが可能のようですが、こっちのほうが貼るだけ、ですよね!)

ukuku09コンテスト(Hamako Online Judge) ghoststudents2

問題文
クエリ2の[c,d]を、二次元平面上の点(c,d)とみなして考えてみましょう。  この時、クエリ1はすべて、「ある長方形範囲に1加算する」ということになります。
この長方形の範囲は、先読み+座標圧縮や、単純な平衡二分木の操作で簡単に求まるので、 問題は$O(Q * log (10^6) * log Q)$で求めることが可能です。 なお、この問題の想定解法である$O(Q * N * log N)$の定数倍高速化よりこちらの方法のほうがはるかに簡単に解くことができますし、理論上はこちらのほうが圧倒的に速いです。
(実際は平衡二分木による定数倍がかなり重く、AVL木や赤黒木などの割柏高速なデータ構造でないとAC出来ませんでした。)

 以上のように、足し算だけでなくxorのような累積可能な演算子、また一次元区間のクエリにこのデータ構造は強いことがわかります。一次元区画のクエリといえば、AtCoder Regular Contest 067のYakiniku Restaurantsもこれで解くことが(理論上)可能です。ただ、あの問題はクエリ制ではないので、imos法でも解けてしまいます。なんでもかんでもデータ構造に頼るとかえってTLEやバグの原因にもなるので注意ですね…

本題2

もう一つ問題を考えてみましょう!

「平面上に点がN個散らばっている。i個目の点の座標は(i,Y_i)であり、コストはC_iである。 今から2種類のクエリを合計Q個処理せよ。
クエリa:(x,y)が与えられ、Y_xにyを代入する
クエリb:(l,r,a)が与えられ、幅[l,r)、高さ[a,∞)の長方形内に存在する点のコストの合計を出力する N,Q≦10^5」

例:
N=6 Y={1,1,4,5,1,4} C={8,10,1,9,1,9}
クエリb,(3,8,3,7)→1+9+9=19
クエリa,(5,5)
クエリa,(3,2)
クエリa,(6,7)
クエリb,(3,8,3,7)→9+1=10

f:id:platypus999:20170920212449j:plain

解法: セグ木にpair<int,int>などの2つの値をもつ平衡二分木を持たせます。 次に、各iについて、[i,i+1)の区間に対応するノードに(Y_i,C_i)を代入します。 また、各区間[l,r)の平衡二分木に、その二つの子[l,m)と[m,r)の平衡二分木の全要素を入れます。 すると、図のようになります。(蟻本の「領域木」参照)

f:id:platypus999:20170920212452j:plain

 空間計算量は$O(N log N)$です。また、構築にかかる時間計算量は$O(N log^2 N)$であり、すこし操作を加える(マージソート的なテクニックを使う)と$O(N log N)$に落とせます。ただ、$log N$が一つ減るだけですので、コストパフォーマンス的におすすめは余りしません。

 クエリaはxにかぶるすべての区間(l≦x<rであるような区間[l,r))について、平衡二分木からY_iを削除し、xを挿入します。計算量はO(log^2 N)です。
また、クエリbは[l,r)を[l1,r1),[l2,r2),[l3,r3),[l4,r4),...,[ls,rs)というs個の区間に分け、それぞれの区間の平衡二分木について、値がa以上の全要素の和を取得し、その第二値の総和を出力すればよいです。 この処理は、ちょうど本題1で説明したアルゴリズムと逆のことをやっています。即ち、 区間更新一点取得か一点更新区間取得か、ということです。 以上のアルゴリズムを使うと、Q個のクエリに対し、$O(Q * log^2 N)$の計算量で処理できます。 $O(N) = O(Q)$の場合$、O(N log^2 N)$となるため、愚直にやった場合の$O(N^2)$に比べて改善がされています。

証明: アルゴリズムの正当性はほぼすべて本題1と被るので割愛いたします。 こちらのほうが区間に対する平衡二分木の取り方がシンプルなので、直感的かもしれません。

実践:

Codeforces Round #433 Div1C Boredom

問題文
すこし考察するとまんま「本題2」の問題であることがわかります。  座標圧縮すらする必要もなく、$O(N log^2 N)$で投げるだけです。 (なお、wavelet matrixというさらに怖いデータ構造でも殴れるらしいです)

Codeforces Round #431 Div1C Goodbye Souvenir

問題文
この問題は、b_i = (同じshapeである最も近くにある左の図形の位置)とすると、  二番目のクエリは「l≦i≦r∧b_i≧lを満たすすべてのi-b_iの和」となります。 これは、すべての1≦i≦Nについて、位置(i,b_i)、コスト(i-b_i)の点があるとき、 幅[l,r)、高さ[l,∞)の長方形内に存在する点のコストの合計であることがわかります。 よって、この問題もセグ木に平衡二分木を載せることによって解くことが可能です。 (この問題はクエリ先読みが可能なので、セグ木にBITを載せることでも解くことができます)


 このように、こっちのデータ構造も一次元区間のクエリにある程度強いことがわかります。 領域木の上位互換であり、かつ実装も(平衡二分木がすでに手元にあるならば)意外と単純なので、こっちも使ってみてはいかがでしょうか?

まとめ

 以上はセグ木に平衡二分木を載せると何ができるかの紹介でした!
確かにライブラリ持ち込み不可の環境では文字数の多さや実装量の多さから余りお勧めはできないかもしれませんが、手元に置いておくといざというときに「貼るだけ」を体現できるようになるはずです!また、領域木やBITを持つセグ木よりも先読みをしなくてよく、煩雑な処理を挟まなくてもよいことが魅力だと思います。みんなも平衡二分木とセグ木を組み合わせて、さらに充実した競プロライフを楽しんでください!

あとでソース貼るかも!