「セグメント木に行列を乗せる」とはどういうことか

この記事は元々ある問題のお題箱リクエストのために書いていたのですが、一般的にテクニックを紹介する記事＋その問題を例題として解説、という流れにすることにしました。

まれによく出題される、「セグメント木に行列を乗せる」とはどういうことか。多くの解説では「モノイドなので乗ります」の一言で片付けられがちですが、そこはまだハードルが高いという人向けに、なるべく代数学の言葉を使わずに「行列とその積を扱う」という点に絞ってまとめます。

「セグメント木」と「行列とその積」については既知であることを前提としますので、リンク先などを参照してください。

セグメント木おさらい

この記事では、

更新クエリ：1点の値を更新する
取得クエリ：指定された区間内の値に対する"何か"を求める

という機能を持つセグメント木だけを扱います。よく使うのは各値が整数型で、区間内の値の総和/最大値/最小値などを取れるものです。もう少し凝った例だと最大公約数なんかを求めることもありますね。

このようなセグメント木を特徴づけるのは、

値の型
区間内の値に対して求める"何か"を定義するような、値同士の二項演算（二引数関数）

です。一応単位元もセットではありますが、これは型と二項演算から自然に決まることが多いです。

ただどんな型と二項演算でもセグメント木で扱えるというわけではなく、いわゆる結合法則と呼ばれる性質を満たす必要があります。例えば $a, b, c$ を整数としたときに、

$a + b + c = (a+b) + c = a + (b + c)$
$\max(a, b, c) = \max(\max(a, b), c) = \max(a, \max(b, c))$

などが成り立ちます。さらに言うと $b+c+d+e+f+g$ みたいな長い式を、どこの足し算から処理しても結果は同じです。これが結合法則です。

セグメント木の計算では、更新クエリでも取得クエリでも「区間の結果と区間の結果をマージして、より大きい区間の結果を得る」というのがキモの動作です。そのため、区間に対して求めるべき値が二項演算の繰り返しで求められ、区間内のどこから計算しても結果が同じになるという性質が必要なのです。

f:id:betrue12:20200330214030p:plain

上記の2例は二項演算 $f(x, y)$ をそれぞれ $f(x, y) = x+y$ , $f(x, y) = \max(x, y)$ としたものです。ここでいわゆる交換法則は必要ではなく、 $f(x, y) = f(y, x)$ は成り立っていなくても大丈夫です（この2例だと成り立ってしまいますが）。

そして単位元は必ず $f(x, e) = f(e, x) = x$ となるような値 $e$ です。足し算なら $0$ 、 $\max$ なら $-\infty$ （実装上は十分小さい整数）です。

これがセグメント木の性質に関する、少し抽象化したおさらいです。ちなみにこれが、よく「モノイド」と言われるものの中身だったりします。

行列を乗せてみる

上記の内容を踏まえれば、話は早いです。行列の積は結合法則を満たします。つまり値の型として行列を、二項演算としてその積を採用すると、これはセグメント木で計算することができます。セグメント木の各ノードが行列を持っていて、更新クエリや取得クエリで区間と区間の値同士をマージするたびに行列積を計算する…ということです。

このセグメントはどのような実装になるでしょうか。骨組みは整数を要素とするセグメント木と変わりません。例えば整数の $\max$ が取れるセグメント木を既に持っているならば、これを書き換えて

各ノードの値の型、単位元の型、更新クエリ関数の引数である値の型、取得クエリ関数の戻り値の型を全て行列（二重配列）に置き換える
更新クエリ関数と取得クエリ関数の中にある値同士の $\max$ 関数を全て行列積の関数に置き換える
単位元を単位行列に置き換える

とすると、行列積が計算できるセグメント木が手に入ります。

ですが型や演算の異なるセグメント木をいちいち書き換えたり、あらかじめ大量に用意するのは面倒です。そのため各プログラミング言語の機能を用いた「抽象化」をする人が多いです。これは具体的に言うと、セグメント木を生成するときのコンストラクタに引数などの形で「値の型」「二項演算の関数」「単位元」を与えることで、様々な種類のセグメント木を作れるような「骨組み」のクラスを作っておく、ということです。

実現手段は言語によって様々ですが、例えばC++であれば

値の型：template機能を用いる
二項演算：コンストラクタにラムダ式で渡して、クラス内ではstd::functionなどで受ける
単位元：コンストラクタに引数として渡す

などの方法があります。

実装例は後でリンクを貼る提出コードを見ていただければと思いますが、コンストラクタを呼び出している部分だけ抜き出すと以下のようになります。（問題に合わせて）この実装では各行列のサイズを $6 \times 6$ 固定であるとしています。

typedef vector<vector<uint32_t>> Mat;
Mat E(6, vector<uint32_t>(6));
for(int i=0; i<6; i++) E[i][i] = 1;

Segtree<Mat> st(N, [](Mat& A, Mat& B)->Mat{
    Mat C(6, vector<uint32_t>(6));
    for(int k=0; k<6; k++) for(int i=0; i<6; i++) for(int j=0; j<6; j++) C[i][j] += B[i][k] * A[k][j];
    return C;
}, E);

テンプレート引数として型名 Mat を与え、引数には順番に要素数 N 、行列積を計算する二項演算のラムダ式、単位元 E を与えています。セグメント木本体のコードには「骨組み」だけがあり、行列積に特有の処理は一切出てきません。これが抽象化の良いところです。

計算量について

32bit/64bit型に収まる整数の四則演算や $\max$ 等は $O(1)$ と見なされることが多いので、そのようなセグメント木のクエリ計算量は要素数を $N$ として $O(\log N)$ と言われます。ただし、例えば行列積などの場合は相応の計算量が掛かるので、クエリ当たりの計算量は $O((マージ1回の計算量) \times \log N)$ になります。サイズ $M$ の正方行列の積であれば、クエリ当たり $O(M^{3}\log N)$ です。

問題例

では実際にどんな問題で使うのかを見てみましょう。元々リクエストの対象だったこの問題を使います。

D - DISCO!

この問題はクエリ $(L, R)$ が1種類であれば以下のようなDPで解けます。

$dp\lbrack i\rbrack\lbrack j\rbrack =$ 文字列を $i-1$ 文字目まで見終わって、DISCO のうち前 $j$ 文字を取る取り方の個数

$dp\lbrack i\rbrack$ から $dp\lbrack i+1\rbrack$ への遷移はこうなります。

$i$ 文字目を採用しない：全ての $j$ について $dp\lbrack i\rbrack\lbrack j\rbrack$ から $dp\lbrack i+1\rbrack\lbrack j\rbrack$ に遷移する
$i$ 文字目を採用する： $S$ の $i$ 文字目が DISCO の $k$ 文字目だとして、 $dp\lbrack i\rbrack\lbrack k-1\rbrack$ から $dp\lbrack i+1\rbrack\lbrack k\rbrack$ に遷移する

クエリ $(L, R)$ について解く場合は、初期条件を $dp\lbrack L\rbrack\lbrack 0\rbrack = 1$ として、DPを計算し $dp\lbrack R+1\rbrack\lbrack 5\rbrack$ を答えとすれば良いです。

ただし実際には大量のクエリに答える必要があるため、ここで行列積セグメント木の出番です。1つの $dp\lbrack i\rbrack$ をベクトルと見なすと、先ほどの遷移は以下のようにベクトルに行列を掛け算することで表現できます。

f:id:betrue12:20200330215759p:plain

クエリ $(L, R)$ の範囲における遷移は、この行列の同区間における積を、 $dp\lbrack L\rbrack\lbrack 0\rbrack = 1$ とした初期条件ベクトル $dp\lbrack L\rbrack$ に掛けることに対応しています。そして得られた $dp\lbrack R+1\rbrack\lbrack 5\rbrack$ が答えなので、行列積の対応する成分を見るとそのまま答えになります。