模擬国内2017G へやわり!(Room Assignment)

前置き

解説の天才二項係数が理解できなかったので,形式的べき級数を用いた機械的(?)な解法の導出について書きます.

問題

onlinejudge.u-aizu.ac.jp

解法

$(i, a _ i)$ の辺を張った無向グラフ $G$ を考える.$G$ から多重辺を取り除いてできる新しいグラフ $G'$ に関して,$G'$ がいくつかのパスの集合となっていなければ valid な部屋割りは存在しないので,答えは $0$ である.以下,$G'$ がパスの集合である場合を考える.

条件「任意の二つの異なる提案された部屋割り A, B に対して,部屋割り A の空き部屋に誰か1人を移しても部屋割り B にならない」について,長さ $3$ 以上のパスに属する人を移動させることで他の valid な部屋割りを得られることはないので,長さ $2$ のパスに属する人を移動させる場合のみを考えればよい.$a _ i \neq i$ より, 長さ $1$ のパスは存在しないことに注意する.

長さ $2$ のパスの個数を $N$, 長さ $3$ 以上のパスの個数を $M $ とおく.$N + M $ 個のパスたちを並べた後に空き部屋を挿入することで部屋割りを得ることを考える.一旦パスの向きを無視して,最後に $2 ^ {N + M}$ を答えに掛けることにする.

例えば [oo][oo] ([oo] は長さ $2$ のパスを表す) に対しては,空き部屋の挿入方法は x[oo][oo], [oo]x[oo], [oo][oo]x の $3$ つある.x[oo][oo][oo]x[oo][oo]x[oo][oo][oo]x は人を $1$ 人選んで動かすことで得られる組合せなので,同時に部屋割りの集合に含めることは出来ない.従って,条件に違反しないようにできるだけ多くの部屋割りを選ぶには,x[oo][oo], [oo][oo]x の $2$ つを選ぶことになる.一般に,[oo] が $i$ 個続く場合,空き部屋の挿入方法は最大で $i + 1 - \lceil i / 2 \rceil$ 個選択することが出来ることが分かる.

さらに,例えば [ooo][oo][oo][ooo][ooo][oo][oo][oo] に空き部屋を挿入することを考えると,[oo] が連続する極大な区間たちに注目することで,この場合に空き部屋の挿入の仕方は最大で $N + M + 1 - \lceil 2 / 2 \rceil - \lceil 3 / 2 \rceil$ 個選択することが出来ることが分かる.

以上の観察を元に,$M $ 個の長さ $3$ 以上のパスを予め並べておき,$M+1$ 個の隙間に合計 $N$ 個の長さ $2$ のパスを挿入することを考える.それぞれの隙間に長さ $2$ のパスが幾つ挿入されるかのみに注目すればよいので,パスが選ばれる順番は区別せずに数えた後に $N! M!$ を答えに掛けることにする.

動的計画法により数えることを考える.

$$\begin{aligned} \mathrm{dp}[i][j][k] = &\text{$i$ 個の隙間に合計 $j$ 個の長さ $2$ のパスを挿入する方法であって,}\\ &\text{空き部屋の挿入方法を最大 $i + j - k$ 個選べるようなものの個数} \end{aligned}$$

とすると,次のように初期化および遷移を書くことが出来る.求めたいのは $S := \displaystyle \sum _ {k} (N+M+1-k)\mathrm{dp}[M+1][N][k]$ である.

$$ \mathrm{dp}[i][j][k] = \begin{cases} 1 & \text{if } i = 0\land j = 0\land k = 0 \\ 0 & \text{if } i = 0\land (j \neq 0 \lor k \neq 0)\\ \displaystyle \sum _ {p = 0} ^ {\infty} \mathrm{dp}[i - 1][j - p][k - \lceil p / 2 \rceil] & \text{otherwise} \end{cases}. $$

従って,形式的べき級数 $\displaystyle f _ {i}(x,y) = \sum _ j \sum _ k \mathrm{dp}[i][j][k] x ^ j y ^ k$ を定義すると,以下が成り立つ.

$$\begin{aligned} f _ 0 &= 1, \\ f _ {i + 1} &= f _ i \cdot (x ^ 0 y ^ 0 + x ^ 1 y ^ 1 + x ^ 2 y ^ 1 + x ^ 3 y ^ 2 + x ^ 4 y ^ 2 + \cdots) \\ &= f _ i \left(\sum _ {i = 0} ^ {\infty} x ^ {2i} y ^ i + xy\sum _ {i = 0} ^ {\infty} x ^ {2i} y ^ i\right) \\ &= f _ i \cdot \dfrac{1 + xy}{1 - x ^ 2 y}. \end{aligned}$$

即ち,$f _ {M + 1} = \left(\dfrac{1 + xy}{1 - x ^ 2 y}\right) ^ {M + 1}$ を得る.

ここで,$S = \displaystyle \sum _ {k} (N+M+1-k) [x ^ N y ^ k] f _ {M + 1}$ が畳込みの形をしていることに注目する.$\displaystyle \sum _ {i} (i + 1) y ^ i = \dfrac{1}{(1 - y) ^ 2}$ であるから,次が成り立つ.

$$ S = [x ^ N y ^ {N + M}] \dfrac{1}{(1 - y) ^ 2}\cdot \left(\dfrac{1 + xy}{1 - x ^ 2 y}\right) ^ {M + 1}. $$

右辺の $x ^ N$ の係数に注目する.

$$\begin{aligned} (1 + xy) ^ {M + 1} &= \sum _ {i = 0} ^ {M + 1}\binom{M + 1}{i} y ^ i \cdot x ^ i, \\ \dfrac{1}{(1 - x ^ 2 y) ^ {M + 1}} &= \sum _ {i = 0} ^ \infty \binom{M + i}{i} y ^ i \cdot x ^ {2i}, \end{aligned}$$

であるから,次を得る.

$$\begin{aligned} S = [y ^ {N + M}] \dfrac{1}{(1 - y) ^ 2} \sum _ {i} \binom{M + 1}{N - 2i} \binom{M + i}{i} y ^ {N - i} \end{aligned}$$

階乗を前計算しておくことで,これは $\Theta(N + M)$ 時間で計算できる.求める答えは $2 ^ {N + M} N! M! S$ であるから,全体 $\Theta(N + M + \log \mathrm{mod})$ で本問題を解くことが出来た.

実装

注: AtCoder Library 使用

#include <iostream>

#include <atcoder/modint>
#include <atcoder/dsu>

using mint = atcoder::modint1000000007;

void remove_multiedges(std::vector<std::vector<int>>& g) {
    std::vector<uint8_t> exists(g.size(), 0);
    for (auto& vs : g) {
        for (int v : vs) exists[v] = true;
        vs.erase(std::remove_if(vs.begin(), vs.end(), [&](int v) { return not std::exchange(exists[v], false); }), vs.end());
    }
}

constexpr int L = 200000;

int main() {
    std::ios::sync_with_stdio(false);
    std::cin.tie(nullptr);

    std::vector<mint> fac(L + 1), fac_inv(L + 1);
    fac[0] = 1;
    for (int i = 1; i <= L; ++i) {
        fac[i] = fac[i - 1] * i;
    }

    fac_inv[L] = fac[L].inv();
    for (int i = L; i > 0; --i) {
        fac_inv[i - 1] = fac_inv[i] * i;
    }

    auto binom = [&fac, &fac_inv](int n, int r) -> mint {
        if (r < 0 or r > n) return 0;
        return fac[n] * fac_inv[r] * fac_inv[n - r];
    };
    
    while (true) {
        int k;
        std::cin >> k;

        if (k == 0) break;

        std::vector<int> a(k);
        for (auto& e : a) std::cin >> e, --e;

        atcoder::dsu uf(k);

        std::vector<std::vector<int>> g(k);
        for (int i = 0; i < k; ++i) {
            g[i].push_back(a[i]);
            g[a[i]].push_back(i);
            uf.merge(i, a[i]);
        }
        remove_multiedges(g);

        int n = 0, m = 0;

        bool all_path = true;
        for (auto group : uf.groups()) {
            n += group.size() == 2;
            m += group.size() >= 3;

            int c1 = 0, c2 = 0;
            for (int v : group) {
                c1 += g[v].size() == 1;
                c2 += g[v].size() == 2;
                all_path &= g[v].size() <= 2;
            }
            all_path &= c1 == 2;
        }

        if (all_path) {
            std::vector<mint> f(n + m + 1);
            for (int i = 0; 2 * i <= n; ++i) {
                f[n - i] = binom(m + i, m) * binom(m + 1, n - 2 * i);
            }

            std::partial_sum(f.begin(), f.end(), f.begin());
            std::partial_sum(f.begin(), f.end(), f.begin());

            std::cout << (mint(2).pow(n + m) * fac[n] * fac[m] * f[n + m]).val() << '\n';
        } else {
            std::cout << 0 << '\n';
        }
    }
    return 0;
}

矩形加算・矩形和取得クエリ

はじめに

次のような問題を考えましょう。

問題 (Rectangle Add Rectangle Sum)

整数 $x, y$ に対して、$2$ 次元平面上の矩形領域 $[x, x + 1) \times [y, y + 1)$ をマス $(x,y)$ とよぶ。初め、全てのマスの重みは $0$ である。整数 $x, y$ に対して、マス $(x,y)$ の重みを $w _ {x, y}$ と表すことにする。

以下のいずれかの形式で表されるクエリが合計で $Q$ 個与えられるので、順番に処理せよ。

  • $\mathrm{add}(v,l,r,d,u)$ : 矩形領域 $[l, r) \times [d, u)$ に含まれる全てのマスの重みに $v$ を加算する。即ち、$l \leq x \lt r$ かつ $d \leq y \lt u$ を満たす全ての整数 $x, y$ について、$w _ {x, y} \leftarrow w _ {x, y} + v$ と更新する。
  • $\mathrm{sum}(l,r,d,u)$ : 矩形領域 $[l, r) \times [d, u)$ に含まれる全てのマスの重みの総和を計算する。即ち、$\displaystyle \sum _ {x = l} ^ {r - 1} \sum _ {y = d} ^ {u - 1} w _ {x, y}$ を計算する。

本記事では、$\mathrm{sum}$ クエリより後に $\mathrm{add}$ クエリが来ない場合の問題 (Static Rectangle Add Rectangle Sum) と、そのような制約が無い一般の問題 (Rectangle Add Rectangle Sum) を扱います。

クエリ先読みが可能という仮定の下で、Static Rectangle Add Rectangle Sum を $\Theta(Q)$ space, $\Theta(Q \log Q)$ time で解くアルゴリズムおよび Rectangle Add Rectangle Sum を $\Theta(Q \log Q)$ space, $\Theta(Q (\log Q) ^ 2)$ time で解くアルゴリズムを解説します。

Static Rectangle Add Rectangle Sum の解法

追記 (2022/05/10) : Rectangle Add Rectangle Sum における解法のようにクエリを 4 分割すれば、最後の遅延セグメント木パートを (通常の) Binary Indexed Tree に代えることができ、定数倍で本解法よりも高速であろうことを確認しました。

$\mathrm{add}$ クエリおよび $\mathrm{sum}$ クエリを次のように分割します。

$\mathrm{add}(v,l,r,d,u)\rightarrow \mathrm{add}(v,l, \infty,d,u), \mathrm{add}(-v,r,\infty,d,u)$

$\mathrm{sum}(l,r,d,u) = \mathrm{sum}(-\infty,r,d,u) - \mathrm{sum}(-\infty,l,d,u)$

この分割により、$\mathrm{add}(v,l, \infty,d,u)$ の形の $\mathrm{add}$ クエリおよび $\mathrm{sum}(-\infty,r,d,u)$ の形の $\mathrm{sum}$ クエリしか存在しないようなケースに帰着できました。ここで、帰着後の問題のクエリ数は、元のクエリ数の定数倍 (ちょうど $2$ 倍) に収まっていることに注意します。

結論を先に書くと、この簡単化された問題は $x$ 方向の平面走査により区間一次関数加算・区間和取得の遅延伝播セグメント木で解くことが出来ます。

以降、$\mathrm{add'}(v,l,d,u)$ で $\mathrm{add}(v,l, \infty,d,u)$ を、$\mathrm{sum'}(r,d,u)$ で $\mathrm{sum}(-\infty,r,d,u)$ を表すことにします。

$\mathrm{sum'}(x,\ast,\ast)$ の計算結果に影響するのは、$l \lt x$ を満たす $\mathrm{add'}(\ast,l,\ast,\ast)$ クエリに限られます。そして、$l \lt x$ のとき、任意の整数 $y$ に対して $\mathrm{add'}(v,l,y,y+1)$ クエリの $\mathrm{sum'}(x,y,y+1)$ への寄与は $v(x-l) = vx - vl$ という $x$ に関する一次関数の形で書くことが出来ます (下図参照)。また、$l = x$ の場合もこの寄与の式は正しいです。

$\mathrm{add'}(v,l,y,y+1)$ の $\mathrm{sum'}(x,y,y+1)$ への寄与

ここで、$\mathrm{add'}(\ast,x,\ast,\ast)$ や $\mathrm{sum'}(x,\ast,\ast)$ を $x$ の昇順に並び替えて処理することにすれば (平面走査)、$\mathrm{sum'}$ クエリに対する答えを保存しつつ、上記の $l \lt x$ (あるいは、$l\leq x$) という制約を外すことができます。先の議論から、タイブレークは任意に定めてよいです。

従って、今までに見た $\mathrm{add'}$ クエリの $\mathrm{sum'}(x,y,y+1)$ への寄与の総和を $w _ y (x)$ と書くことにすれば、各クエリは次のように処理できます。

  1. $\mathrm{add'}(v,l,d,u)$ の場合:

    全ての $l \leq y \lt d$ を満たす整数 $y$ に対して $w _ y(x) \leftarrow w _ y(x) + vx - vl$ と更新する。

  2. $\mathrm{sum'}(r,d,u)$ の場合:

    $\displaystyle \sum _ {y = d} ^ {u - 1} (w _ y(r)) = \left(\sum _ {y = d} ^ {u - 1} w _ y \right) (r)$ を計算する。

これらの操作は、予めクエリに現れる $y$ 座標を列挙して座標圧縮を行っておくことにより、区間一次関数加算・区間和取得の遅延伝播セグメント木により各々 $\Theta(\log Q)$ time で処理できます。

ここで、区間一次関数加算は $(ax + b) + (cx + d) = (a + c)x + (b + d)$ のような操作を指し、区間等差数列加算とは異なるものです。混同しないよう注意してください。

前処理の座標圧縮やクエリのソートなどは $\Theta(Q)$ space, $\Theta(Q \log Q)$ time で処理できるので、全体 $\Theta(Q)$ space, $\Theta(Q \log Q)$ time で Static Rectangle Add Rectangle Sum を解くことが出来ました。

実装

私のライブラリへのリンクで代えさせて頂きます。(追記 (2022/05/10) : より定数倍のよい Binary Indexed Tree を用いた実装に差し替えました)

suisen-cp.github.io

Rectangle Add Rectangle Sum の解法

$\mathrm{add}$ や $\mathrm{sum}$ が入り混じっている場合、先の解法のようにクエリをうまく並び替えることは難しいです。

以下のようにクエリをさらに分割し、さらに計算しやすい形にします。

$\mathrm{add}(v,l,r,d,u)\rightarrow \mathrm{add}(v,l,\infty,d,\infty), \mathrm{add}(-v,l,\infty,u,\infty), \mathrm{add}(-v,r,\infty,d,\infty), \mathrm{add}(v,r,\infty,u,\infty)$

$\mathrm{sum}(l,r,d,u)=\mathrm{sum}(-\infty,r,-\infty,u)-\mathrm{sum}(-\infty,r,-\infty,d)-\mathrm{sum}(-\infty,l,-\infty,u)+\mathrm{sum}(-\infty,l,-\infty,d)$

この分割により、$\mathrm{add}(v,l, \infty,d,\infty)$ の形の $\mathrm{add}$ クエリおよび $\mathrm{sum}(-\infty,r,-\infty,u)$ の形の $\mathrm{sum}$ クエリしか存在しないようなケースに帰着できました。ここで、帰着後の問題のクエリ数は、元のクエリ数の定数倍 (ちょうど $4$ 倍) に収まっていることに注意します。

以降、$\mathrm{add''}(v,l,d)$ で $\mathrm{add}(v,l, \infty,d,\infty)$ を、$\mathrm{sum''}(r,u)$ で $\mathrm{sum}(-\infty,r,-\infty,u)$ を表すことにします。

$\mathrm{sum''}(x,y)$ の計算結果に影響するのは、$l \lt x$ かつ $d \lt y$ を満たす $\mathrm{add''}(\ast,l,d)$ クエリに限られます。そして、$l \lt x$ かつ $d \lt y$ のとき、$\mathrm{add''}(v,l,d)$ クエリの $\mathrm{sum''}(x,y)$ への寄与は $v(x-l)(y-d) = vxy - vdx-vly+vld$ であり、即ち $f(x,y)=axy+bx+cy+d$ という形の関数で書くことが出来ます (下図参照)。

$\mathrm{add''}(l,d)$ の $\mathrm{sum''}(x,y)$ への寄与

従って、この簡単化された問題は $axy+bx+cy+d$ という形をした $x,y$ を変数とする関数を点の重みとする Point Add Rectangle Sum に帰着することが出来ます。

Point Add Rectangle Sum は例えば必要なところだけ作る 2D BIT などを用いることで $\Theta(Q \log Q)$ space, $\Theta(Q (\log Q) ^ 2)$ time で解くことが出来るので、Rectangle Add Rectangle Sum も同じ計算量オーダーで解くことが出来ました。

ところで、Point Add Rectangle Sum における重みの計算コストが通常の整数の加算の場合と比較して単純計算で $4$ 倍であり、クエリ数も $4$ 倍なので、定数倍はかなり大きいです。あとに示す実装例で $10 ^ 5$ クエリのランダムケースで計測すると $1$ 秒以上掛かっていました。

実装

私のライブラリへのリンクで代えさせて頂きます。

suisen-cp.github.io

おわりに

Rectangle Add Rectangle Sum の方は Point Add Rectangle Sum に帰着してしまえば最悪 Library Checker から借りられるけど、Static の方は予めライブラリを作っておかないとかなり苦しい気がする... (Static の方もクエリを $4$ 分割の解法なら割と軽めでした)

そらで書ける <Θ(N), Θ(log N)> Level Ancestor

次の問題を考えます.

問題 (Level Ancestor Problem)

頂点 $1$ を根とする $N$ 頂点の木 $T$ が与えられます.

頂点 $v$ と非負整数 $k$ に対して,$\mathrm{LA}(v, k)$ を頂点 $v$ から親方向への辺をちょうど $k$ 回辿って到達する頂点と定義します.

頂点と非負整数の組が $Q$ 個与えられます.$i$ 番目の組を $(v _ i, k _ i)$ として,各 $i$ に対して $\mathrm{LA}(v _ i, k _ i)$ を求めて出力してください.ただし,頂点 $v _ i$ から親方向への辺を $k _ i$ 回辿れない場合は,代わりに $-1$ を出力してください.

本記事の意義 (?)

本節は読み飛ばして貰っても構いません.

Level Ancestor Problem の解法として,Heavy Light Decomposition (HLD) を用いたものや,ダブリングによるものがよく用いられています (要出典).それぞれの解法について詳しくは述べませんが,計算量は次のようになっています.

前計算 クエリ 空間
HLD $\Theta(N)$ $\Theta(\log N)$ $\Theta(N)$
ダブリング $\Theta(N\log N)$ $\Theta(\log N)$ $\Theta(N \log N)$
本記事で紹介するアルゴリズム $\Theta(N)$ $\Theta(\log N)$ $\Theta(N)$

上表にも示しましたが,本記事では Level Ancestor Problem を前計算 $\Theta(N)$,クエリ $\Theta(\log N)$ で解くアルゴリズムを紹介します.表を見ると HLD を用いた解法で十分そうに見えます.実際 HLD で十分なのですが,Level Ancestor Problem を解くためだけに HLD を実装するのは,やや大変です.実装の簡単さの点ではダブリングによる解法は優れていますが,前計算に掛かる時間計算量や空間計算量のオーダーが HLD のそれと比較して悪化してしまいます.本記事では,HLD による解法と同じ計算量オーダーを達成しつつ,実装も簡単なアルゴリズムを紹介します.

なお,クエリ先読みが可能な場合 (Offline Level Ancestor Problem) は,以下の記事で紹介されている手法で解くのが楽だと思います (計算量の面でも,優れています).

noshi91.hatenablog.com

本題

書き切ってから気付いたんですが,先に「おわりに」の節を読んだ方が何をしようとしているかわかりやすいかもしれません.


クエリ v k の処理を考えましょう.頂点 $v$ の深さを $\mathrm{de p}(v)$ とすると,$\mathrm{LA}(v, k)$ が存在するための必要十分条件は $\mathrm{de p}(v) \geq k$ を満たすことです.そこで,以降は $\mathrm{de p}(v) \geq k$ を仮定します.

深さがちょうど $d$ であるような頂点の集合を $S _ d$ とすると,$\mathrm{LA}(v, k) \in S _ {\mathrm{de p}(v) - k}$ が成り立ちます.また,$S _ {\mathrm{de p}(v) - k}$ に含まれる頂点のうち,$v$ の先祖 (ここでは,$v$ 自身も含みます) であるようなものはちょうど $1$ つしか存在しないので,$S _ {\mathrm{de p}(v) - k}$ に含まれる $v$ の先祖を高速に見つけられればよいです.しかし,$S _ {\mathrm{de p}(v) - k}$ のサイズは最大で $\Theta(N)$ となるため,線形探索では最悪の場合は $\Omega(N)$ time 掛かってしまいます.そこで,上手く探索出来るように工夫をします.

木 $T$ に対して根 $1$ を始点とする深さ優先探索を行い,その訪問順 (pre-order) に頂点を並べ替えてできる頂点の列を $v _ 1, \ldots, v _ N$ とします.そして,$v _ i$ の頂点番号が $i$ となるように頂点番号を振り直します.古い頂点番号から新しい頂点番号への変換やその逆変換は各順列を覚えておくことで定数時間で行えます.そこで,以降は新しく振り直した頂点番号を用いて考えます.

頂点番号の振り直しの例

さて,このような頂点番号の振り方をすることで,頂点 $u$ を根とする部分木のサイズを $\mathrm{sub}(u)$ とすれば,頂点 $u$ を根とする部分木は半開区間 $[u,u+\mathrm{sub}(u))$ を成します.この事実から,$\mathrm{LA}(v,k)$ を根とする部分木が成す区間を考えることで, $\mathrm{LA}(v,k) \leq v$ を得ます.

ここで,$S _ {\mathrm{de p}(v) - k}$ の要素 $x$ であって $x \leq v$ を満たす頂点のうち最大のものを $p$ とすれば,実は $\mathrm{LA}(v,k) = p$ が成り立ちます $(\star)$.従って,$S _ {\mathrm{de p}(v) - k}$ の要素を昇順に並べた列を予め計算しておくことで,二分探索により $\mathrm{LA}(v,k)$ を $\Theta(\log N)$ time で計算することが出来ます.

$(\star)$ の証明

まず,$\mathrm{LA}(v,k) \leq v$ より,$p$ は必ず存在します.

$p$ 以外に $x \leq v$ を満たす頂点 $x \in S _ {\mathrm{de p}(v) - k} \backslash \{p\}$ が存在しない場合は,明らかに $\mathrm{LA}(v,k) = p$ が成り立ちます.

$p$ 以外に $x \leq v$ を満たす頂点 $x \in S _ {\mathrm{de p}(v) - k} \backslash \{p\}$ が存在する場合を考えます.$x \leq v$ を満たす頂点 $x \in S _ {\mathrm{de p}(v) - k} \backslash \{p\}$ を任意に取り,$x = \mathrm{LA}(v,k)$ を仮定します.このとき,$x$ は $v$ の先祖なので,$v \in [x, x + \mathrm{sub}(x))$ が成り立ちます.ところで $p$ は $x$ を根とする部分木には含まれず,また $p$ の最大性から $x\lt p$ が成り立つため,$p \geq x + \mathrm{sub}(x)$ を満たす必要があります.従って,$v \in [x, x + \mathrm{sub}(x))$ より $x \leq v \lt x + \mathrm{sub}(x) \leq p$,即ち $v\lt p$ を得,これは $p \leq v$ であることに矛盾します.

以上より,$p = \mathrm{LA}(v, k)$ が成り立ちます.(証明終)

各 $d$ に対して $S _ d$ を昇順に並べた列 $A _ d$ を計算する必要がありますが,これは深さ優先探索において頂点 $u$ を訪問したタイミングで $A _ {\mathrm{de p}(u)}$ の末尾に $u$ を追加することで,前計算に掛かる計算量を $\Theta(N)$ に保ったまま構築することが出来ます.

以上より,Level Ancestor Problem を前計算 $\Theta(N)$ time,クエリ $\Theta(\log N)$ time,空間 $\Theta(N)$ で解くことが出来ました.

実装

折りたたみ

#include <algorithm>
#include <vector>

struct LevelAncestor {
    LevelAncestor() = default;
    // 隣接リスト g と根 root 
    LevelAncestor(const std::vector<std::vector<int>> &g, const int root = 0) : _n(g.size()), _visit_time(_n), _depth(_n), _bucket(_n) {
        build(g, root);
    }
    
    // u から親を k 回辿って到達する頂点を返す.0<=k<=dep[u] を満たさない場合は,そのような頂点は存在しないので -1 を返す.
    int query(const int u, const int k) const {
        const int d = _depth[u] - k;
        if (d < 0 or d > _depth[u]) return -1;
        // 頂点 i と頂点 j を訪問時間で比較する関数
        auto comp = [this](int i, int j) {
            return _visit_time[i] < _visit_time[j];
        };
        // u 以下のうち最大 = u より真に大きい頂点のうち最小のものの 1 つ前
        // なので,upper_bound をしてから prev で 1 つ戻せばよい
        return *std::prev(std::upper_bound(_bucket[d].begin(), _bucket[d].end(), u, comp));
    }
    int operator()(const int u, const int k) const {
        return query(u, k);
    }

private:
    // 頂点数
    int _n;
    // 深さ優先探索における訪問時刻
    std::vector<int> _visit_time;
    // 頂点の深さ
    std::vector<int> _depth;
    // 記事中の S.各々訪問時刻の昇順にソートされている
    std::vector<std::vector<int>> _bucket;

    void build(const std::vector<std::vector<int>> &g, const int root) {
        // 現在時刻
        int time = 0;
        // u : 現在見ている頂点
        // p : u の親
        auto dfs = [&](auto dfs, int u, int p) -> void {
            // 頂点 u の訪問時刻を記録して,時刻を 1 つ進める
            _visit_time[u] = time++;
            // S _ {dep(u)} に u を追加
            _bucket[_depth[u]].push_back(u);
            for (int v : g[u]) {
                if (v == p) continue;
                _depth[v] = _depth[u] + 1;
                dfs(dfs, v, u);
            }
        };
        dfs(dfs, root, -1);
    }
};

(2022/06/08 追記)

上記の実装は std::upper_bound に渡した比較器内で配列アクセスを行っているため,遅いです.訪問時刻をバケットに詰めて配列アクセスの回数を減らすことで定数倍高速化を図ることが出来ます.

以下に高速な実装例を示しておきます.(そらで書けると謳っておきながら,そこそこ面倒くさい感じになってしまいました)

suisen-cp.github.io

(追記終)

検証

Library Checker

Level Ancestor を verify できる問題が思い浮かばなかったので,LA で二分探索することで LCA を計算するコードを書いて,LCA を verify する問題に投げています.

おわりに

正当性を示すために長々と書いてしまいましたが,アルゴリズムの中身はかなり直感的で実装もしやすいと思います.やっていることは,前計算では DFS の訪問順に深さ毎に用意したバケツに頂点を詰め込んで,クエリでは対応するバケツを見て訪問時刻が自身の訪問時刻を超えないもののうち最も遅く訪れた頂点を二分探索で求めるだけなので.

AtCoder Beginner Contest 245 Ex - Product Modulo 2

問題

atcoder.jp

解法

$M = \prod _ {t = 1} ^ l {p _ t} ^ {q _ t}$ と素因数分解されるとする.中国剰余定理より,各 $t = 1, \ldots, l$ に対して

$$\prod _ {i = 1} ^ K A _ i ^ t \equiv N \pmod {{p _ t} ^ {q _ t}}$$

を満たす列 $\{ A _ i ^ t \} _ {i = 1} ^ K \in (\Z/ {p _ t} ^ {q _ t} \Z) ^ K$ の個数 $x _ t$ を求めれば,答えは $\prod _ {t = 1} ^ l x _ t$ として求めることが出来る.

従って,$t$ 毎に独立に問題を解いてよいので,以降は次の部分問題を考える.

列 $\{ A _ i \} _ {i = 1} ^ K \in (\Z/ p ^ q \Z) ^ K$ であって,$\prod _ {i = 1} ^ K A _ i \equiv N \pmod {p ^ q}$ を満たすものの個数を求めよ.

整数 $x$ に対して,$v _ p(x)$ を $p$ が $x$ を割り切る回数と定義する.便宜上,$v _ p (0) = \infty$ と定義する.このとき,次の事実が成立する.

$X \equiv Y \pmod{ p ^ q }$ ならば,

  • $Y \equiv 0 \pmod {p ^ q}$ のとき,$v _ p (X) \geq q$
  • $Y \not\equiv 0 \pmod {p ^ q}$ のとき,$v _ p (X) = v _ p(Y)$

証明

  • $Y \equiv 0 \pmod {p ^ q}$ のとき

    $X \equiv Y \equiv 0 \pmod {p ^ q}$ より,ある整数 $k$ が存在して $X = k \cdot p ^ q$ を満たすので $v _ p (X) \geq q$.

  • $Y \not \equiv 0 \pmod {p ^ q}$ のとき

    $Y \not \equiv 0 \pmod {p ^ q}$ より $v _ p(Y) \lt q$ であり,$p$ と互いに素な整数 $Y'$ が存在して $Y = Y' \cdot p ^ {v _ p (Y)}$ と書ける.$X \equiv Y\pmod {p ^ q}$ より,ある整数 $k$ が存在して $X = Y + k \cdot p ^ q = p ^ {v _ p (Y)} ( Y' + k \cdot p ^ {q - v _ p (Y)})$ を満たす.

    ここで,$v _ p(Y) \lt q$ より $k \cdot p ^ {q - v _ p (Y)}$ は $p$ の倍数であり,かつ $Y'$ は $p$ と互いに素であるから $Y' + k \cdot p ^ {q - v _ p (Y)}$ も $p$ と互いに素.従って,$X$ が $p$ で割り切れる回数はちょうど $v _ p (Y)$ である.

上記の事実に基づき,$N \equiv 0 \pmod {p ^ q}$ かどうかで場合分けをする.

$N \not\equiv 0 \pmod {p ^ q}$ の場合

$v _ p (\prod _ {i = 1} ^ K A _ i) = \sum _ {i = 1} ^ K v _ p (A _ i) = v _ p (N)$ を満たす必要がある.ここで,$v _ p (N) \lt q$ より $v _ p (A _ i) \lt q$ であることを注意しておく.

非負整数 $x \leq q$ を固定して $p$ と互いに素な整数 $a$ を動かしたときに,$a \cdot p ^ {x} \bmod p ^ q$ は何種類の値を取り得るかを考える.$x \leq q$ の下では,$a \cdot p ^ x \equiv b \cdot p ^ x \pmod {p ^ q}$ であることの必要十分条件は $a\equiv b \pmod {p ^ {q - x}}$ を満たすことであるから,結局 $\varphi (p ^ {q - x})$ 通りの値を取ることが分かる.

さて,各 $v _ p (A _ i)$ を固定した上で,$A _ 1, \ldots, A _ {K - 1}$ を自由に決めてしまうことにする.いま,$v _ p (A _ i) \lt q$ であるから,$A _ 1, \ldots, A _ {K - 1}$ の決め方は

$$\begin{aligned} \prod _ {i = 1} ^ {K - 1} \varphi (p ^ {q - v _ p (A _ i)}) &= \prod _ {i = 1} ^ {K - 1} p ^ {q - v _ p (A _ i) - 1} (p - 1) \\ &= \bigl( p ^ {q - 1} (p - 1) \bigr) ^ {K - 1} \cdot p ^ {-\sum _ {i = 1} ^ {K - 1} v _ p (A _ i)} \\ &= \bigl( p ^ {q - 1} (p - 1) \bigr) ^ {K - 1} \cdot p ^ {-(v _ p (N) - v _ p (A _ K))} \\ \end{aligned}$$

通りである.最後の行では $\sum _ {i = 1} ^ K v _ p (A _ i) = v _ p (N)$ を用いた.

ここで,$p$ と互いに素な整数 $a,b,c$ を用いて $\prod _ {i = 1} ^ {K - 1} A _ i = a \cdot p ^ {v _ p (N) - v _ p (A _ K)}$,$A _ K = b \cdot p ^ {v _ p (A _ K)}$,$N = c \cdot p ^ {v _ p (N)}$と表される.$\prod _ {i = 1} ^ K A _ i \equiv N \pmod {p ^ q} \iff ab\equiv c \pmod {p ^ {q - v _ p(N)}}$ であり,$a$ と $p$ は互いに素であるから,$b$ は $0 \leq b \lt p ^ {q - v _ p(N)}$ の範囲で一意存在する.先の議論から,$A _ K$ の値を考える上で $b$ は $\mathrm{mod}\; p ^ {q - v _ p(A _ K)}$ で同一視されるから,条件を満たす $A _ K$ は $p ^ {v _ p (N) - v _ p (A _ K)} $ 通りだけ存在する.

以上より,$v _ p (A _ i)$ を固定した場合に,条件を満たす $A$ は $\bigl( p ^ {q - 1} (p - 1) \bigr) ^ {K - 1}$ 個存在する.これは $v _ p (A _ i)$ に依らないので,結局全体では $H(v _ p (N), K) \cdot \bigl( p ^ {q - 1} (p - 1) \bigr) ^ {K - 1}$ 個存在する.ここで,$H$ は重複組合せを表し,$\displaystyle H(v _ p (N), K) = \binom{v _ p (N) + K - 1}{v _ p(N)}$ である.

$N \equiv 0 \pmod {p ^ q}$ の場合

このとき,$\prod _ {i = 1} ^ K \equiv N \pmod {p ^ q}$ であるための必要十分条件は,$\sum _ {i = 1} ^ K v _ p (A _ i) \geq q$ を満たすことである.$A _ i \equiv 0 \pmod {p ^ q}$ の場合を考慮する必要があるため,$\varphi (p ^ {q - v _ p(A _ i)})$ の計算に場合分けが発生する.そこで,$A _ i \not\equiv 0 \pmod {p ^ q}$ を満たす $i$ の個数 $j$ を固定して,$j$ の値毎に数え上げることを考える.

  • $j \lt K$ の場合

    このとき,$\sum _ {i = 1} ^ K v _ p (A _ i) \geq q$ は必ず成立する.従って,$j$ を固定したときの求める列の個数は $\displaystyle \binom{K}{j} \cdot \bigl( p ^ {q - 1} (p - 1) \bigr) ^ j \cdot \Biggl(\sum _ {l = 0} ^ {q - 1} p ^ {-l}\Biggr) ^ j = \binom{K}{j} \cdot (p ^ q - 1) ^ j$ である.

    これを $j = 0, \ldots, K - 1$ の範囲で和を取ると,以下のようになる.

    $$\sum _ {j = 0} ^ {K - 1} \binom{K}{j} \cdot (p ^ q - 1) ^ j = (p ^ q) ^K - (p ^ q - 1) ^ K$$

  • $j = K$ の場合

    このとき,$\sum _ {i = 1} ^ K v _ p (A _ i) \geq q$ は必ずしも成立しない.$\sum _ {i = 1} ^ K v _ p (A _ i) = k$ となる列 $A$ の個数は,多項式 $\displaystyle f(x) = \bigl( p ^ {q - 1} (p - 1) \bigr) ^ K \cdot \Biggl(\sum _ {l = 0} ^ {q - 1} p ^ {-l} x ^ l \Biggr) ^ K = (p - 1) ^ K \cdot \Biggl(\sum _ {l = 0} ^ {q - 1} p ^ {q - 1 - l} x ^ l \Biggr) ^ K$ の $x ^ k$ の係数として計算される.

    求めたいのは $x ^ q, x ^ {q + 1}, \ldots $ の係数の和であるが,次数は最大でも $K (q - 1)$ にもなるため,実行時間制限に間に合わせるためには工夫する必要がある.$q$ は最大でも $\lfloor \log _ 2 M \rfloor$ 程度であるから,全ての係数の和から,$x ^ 0, \ldots, x ^ {q - 1}$ の係数の和を差し引くことを考える.全ての係数の和は $f(1)$ として計算され,即ち

    $$(p - 1) ^ K \cdot \Biggl(\sum _ {l = 0} ^ {q - 1} p ^ {q - 1 - l} \Biggr) ^ K = (p - 1) ^ K \cdot \left(\dfrac{p ^ q - 1}{p - 1}\right) ^ K = (p ^ q - 1) ^ K$$

    である.$f(x) \bmod x ^ q$ に関しては,例えば二分累乗法で $O(q ^ 2 \log K)$ や $O(q \log q \log K)$ などで計算することが出来る.

おわりに

難しいよ...

木上の等高線集約クエリ

(追記 2022/03/27 23:07)

変種 1, 1' に対する可逆性を仮定しない解法を教えて頂きました.詳細は以下の記事を参照してください.

noshi91.hatenablog.com

www.mathenachia.blog

(追記終)

次のような問題を考えます.

問題

$N$ 頂点の木 $T$ が与えられる.$v\in\{1,\ldots, N\}$ に対して,頂点 $v$ の重みは $W _ v$ で初期化されている.以下の形式で表されるクエリが $Q$ 個与えられるので,順番に処理せよ.

  • 1 v x : 頂点 $v$ の重みを $x$ に更新する.即ち,$W _ v \leftarrow x$ とする.
  • 2 v d : 頂点 $v$ からの最短距離が $d$ であるような頂点の集合を $S(v,d)$ として,$\displaystyle \sum _ {u \in S(v,d)} W _ u$ を出力する.

ここで,頂点 $u$ と頂点 $v$ の最短距離を木 $T$ 上の $u, v$ パスの辺数の最小値と定義します.

観察

例として以下の図のような木を考えます.また,頂点の重み $W _ i$ は,理解の簡単のため $W _ i = i$ としておきます.

木 $T$

集約クエリを容易に計算できる例として,根に対するクエリが考えられます.以下の図は,頂点 $v=0$ から距離 $d=2$ の頂点を集約するクエリを表しています.赤の頂点が $v$ (以降,これを原点と呼びます),青の頂点は原点 $0$ から距離 $2$ の集約対象の頂点を表しています.このとき,出力すべき値は $W _ 2 + W _ 3 + W _ 9 + W _ {10} + W _ {17} + W _ {18} + W _ {20} = 2+3+9+10+17+18+20=79$ となります.

頂点 $0$ から距離 $2$ の頂点を集約するクエリ

この場合は,幅優先探索の訪問順に頂点番号を振り直すことで区間和を求めるクエリに帰着されます.従って,原点として根しか選ばれない場合は,頂点の重みの和を Segment Tree などを用いて管理することで,タイプ 1 , 2 のクエリを各 $O(\log N)$ time で処理することが出来ます.

難しいのは,原点として根以外が選択される場合です.以下の図は,原点 $1$ から距離 $2$ の頂点を集約するクエリを表しています.このとき,出力すべき値は $W _ 4 + W _ 5 + W _ 8 + W _ {16} = 4 + 5 + 8 + 16 = 33$ となります.

頂点 $1$ から距離 $2$ の頂点を集約するクエリ

この場合,先述の方法で頂点番号を振り直したとしても,最悪の場合は $\Theta(\sqrt{N})$ 個の区間に分割されます $(\star 1)$.最悪ケースの一例としては,以下のようなケースが考えられます.

$\Theta(\sqrt{N})$ 個の区間に分割されるケース

$(\star 1)$ の補足 原点 $v$ から距離 $d$ の頂点を集約するクエリを考えます.

$v$ から $k\; (0\leq k \leq d)$ 回親を辿ったあと,$v$ が属さない部分木に $d - k$ 回潜ることで到達できる頂点は集約対象の頂点であり,その深さは $\mathrm{de p}(v) + d - 2 k$ と表されます.逆に,全ての集約対象の頂点はこれで尽くされます.

親方向に辿る頂点は異なる $k$ で共有できますが,部分木に潜る際に辿る頂点は異なる $k$ で共有できません.従って,集約対象の頂点たちが $l$ 種類の深さを持つとすれば,頂点は少なくとも $0 + 1 + \cdots + (l - 1) = l (l - 1) / 2$ 個必要であり,集約対象となる頂点の深さの種類数は $O(\sqrt{N})$ であることが示されます.各深さに対して区間は高々 2 つなので,上界を $O(\sqrt{N})$ で抑えられることが示されます.

下界に関しては, の構築を一般化することで示されます.(補足終)

本記事では,上記の方法 1 よりも高速にクエリを処理することを目指します.

解法

前計算

唐突ですが,重心分解を用いて本問題を解くことを考えます.重心分解の過程を表す木 $G _ T$ を作ります.$G _ T$ は,次の手続きによって生成される根付き木です.

  1. $T$ の重心の 1 つを任意に選んで $c$ とする.
    • $T$ がただ 1 つの頂点からなる場合は,頂点 $c$ のみからなる,$c$ を根とする根付き木を返す.
    • $T$ が 2 つ以上の頂点からなる場合は,$T$ から頂点 $c$ と $c$ に接続する全ての辺を削除して出来る森の各連結成分ごとに再帰的に本手続きを行う.その結果返された根付き木 $X _ 1, \ldots, X _ k$ を頂点 $c$ の子する,$c$ を根とする根付き木を返す.

「観察」の節で例示した木 $T$ に対しては,例えば以下のように $G _ T$ は構築されます (重心が 2 つある場合の選択には任意性があるため,必ずしも図の通りになるとは限りません).

$G _ T$ の構築 (1)

追記 (2023/11/22) 次の図以降、誤って頂点 $5$ と $7$ を逆に描いてしまっています。また、頂点 $12$ と $15$ も逆に描いてしまっています。画像データを紛失したので、ここで注記するだけに留めさせてください。すみません。 (追記終)

$G _ T$ の構築 (2)

$G _ T$ の構築 (3)

$G _ T$ の構築 (4)

$G _ T$ の構築 (5)

$G _ T$

上記の手続きにおいて,$T$ のすべての頂点はちょうど $1$ 回ずつ重心として選択されます.即ち,$G _ T$ の頂点集合 $V(G _ T)$ と $V(T)$ は一致します.ここで,頂点 $v$ が重心として選択された時点の $v$ を含む連結成分は木となります.これを $v$ を根とする根付き木とし,$U _ v$ で表します.

各 $v$ に対して,$V(U _ v)$ を $U _ v$ における幅優先探索の訪問順に並べ替え,頂点の重みの和を管理する Segment Tree $\mathrm{Seg} _ v$ を構築しておきます.「観察」の節で述べたように,この前計算により $U _ v$ に対する原点を $v$ とする集約クエリおよび重みの更新クエリを効率的に処理することが出来ます.

さて,$G _ T$ の高さは $O(\log N)$ と評価されること $(\star 2)$,および $\displaystyle \sum _ { v = 0 } ^ { N - 1} | V(U _ v) | = O(N \log N)$ が成立すること $(\star 3)$ などから,ここまでの前計算は全体 $O(N \log N)$ time で可能です.

$(\star 2)$ の補足 $G _ T$ における $v$ の親を $p$ とすれば,重心の定義から $|V(U _ v)| \leq \dfrac{1}{2} |V(U _ p)|$ が成立します.$G _ T$ の根 $r$ について $|V(U _ r)| = N$ であること,および任意の頂点 $v$ に対して $|V(U _ v)| \geq 1$ であることから,$G _ T$ の高さが $O(\log N)$ であることが従います.(補足終)

$(\star 3)$ の補足 $G _ T$ において深さが等しい任意の相異なる頂点 $u, v$ に対して $V(U _ u) \cap V(U _ v) = \empty$ が成り立ちます.従って,$G _ T$ において深さが $d$ であるような頂点の集合を $S _ d$ とすれば,任意の $d$ に対して $\sum _ { v \in S _ d } | V(U _ v) | \leq N$ が成り立ちます.$G _ T$ の深さの最大値は $O(\log N)$ であったので,$\sum _ {v = 0} ^ {N - 1} | V(U _ v) | = O(N \log N)$ と評価されます.(補足終)

クエリ処理

更新クエリ

頂点 $v$ に対する更新クエリを考えます.全ての Segment Tree に正しい情報を持たせておくためには,$v \in V(U _ w)$ を満たす全ての頂点 $w$ に対して $\mathrm{Seg} _ w$ を更新する必要があります.$G _ T$ の定義から,$v \in V(U _ w)$ であるための必要十分条件は $G _ T$ において $v$ から親を $0$ 回以上辿ることで $w$ に到達可能であることです.$G _ T$ の高さは $O(\log N)$ であったので,そのような $w$ は $O(\log N)$ 個しかありません.$\mathrm{Seg} _ w$ の更新は $O(\log N)$ time で可能なので,全体 $O( (\log N) ^ 2)$ time で更新クエリを処理することが出来ました.なお,$\mathrm{Seg} _ w$ において頂点 $v$ の値が格納されている添字を高速に取得できる必要がありますが,これは前計算の計算量を悪化させることなく可能です.

集約クエリ

頂点 $v$ から距離 $d$ の頂点の重みの総和 $\displaystyle \sum _ {u \in S(v, d)} W _ u$ を求めるクエリを考えます.

まず,「観察」の節で述べたように,$\displaystyle \sum _ {u \in S(v, d) \cap V(U _ v)} W _ u$ は $\mathrm{Seg} _ v$ に対する区間和取得クエリとなり容易に計算されるので,答えに足し込んでおきます.

続いて,$G _ T$ における $v$ の親 $p$ に対して,$\displaystyle \sum _ {u \in (S(v, d) - V(U _ v)) \cap V(U _ p)} W _ u$ を計算することを考えます.このとき,$U _ p$ における $p$ の子であって,$v$ と一致するか,あるいは $v$ を子孫に持つような頂点が存在するのでこれを $c$ とします.$c$ を根とする子部分木の根を $v$ に変更すれば,$U _ v$ と一致します.即ち,$c$ を根とする子部分木に含まれる頂点であって,$v$ からの距離が $d$ であるような頂点の重みの和は既に答えに足し込まれています.

従って,下図のように,$\displaystyle \sum _ {u \in (S(v, d) - V(U _ v)) \cap V(U _ p)} W _ u$ は $\mathrm{Seg} _ p$ における高々 2 つの区間の和として計算することが出来ます.

$U _ p$

適切な前計算により $c$ や $\mathrm{Seg} _ p$ で和を取るべき区間,$p$ と $v$ の距離 $\mathrm{dist}(v, p)$ は $O(1)$ time で求めることができます $(\star 4)$.従って,$\displaystyle \sum _ {u \in (S(v, d) - V(U _ v)) \cap V(U _ p)} W _ u$ も $O(\log N)$ time で計算することが出来ました.

$(\star 4)$ の補足 例えば,LA や LCA を前計算 $O(N)$ time,クエリ $O(1)$ time で求めるアルゴリズムが存在するので,それらを用いることで前計算の計算量を悪化させること無く達成されます.

あるいは,集約クエリの計算中に現れる $p,v$ の組は $O(N \log N)$ 通りであることを用いて全て前計算することによっても達成可能です.後者の方法は特別なライブラリを必要としないこと,クエリが配列へのアクセスだけで済むことが利点ですが,前計算がやや重くなることが欠点です.

さらに別の選択肢として,$O(\log N)$ time を許容して HLD などで対応しても Segment Tree で区間和を計算するのに結局 $O(\log N)$ time だけ掛かるため集約クエリの計算量は悪化しません.好みに合わせて実装してください.(補足終)

同様の計算を $G _ T$ の根まで遡りながら行うことで,クエリの答えを計算することが出来ます.$G _ T$ の高さは $O(\log N)$ であったので,全体 $O( (\log N) ^ 2)$ time で集約クエリを処理することが出来ました.

実装

私のライブラリへのリンクを張っておきます.

https://suisen-cp.github.io/cp-library-cpp/library/tree/contour_aggregation_query_on_tree.hppsuisen-cp.github.io

問題例

C - 木の問題

全ての頂点の重みを $1$ として集約クエリを投げるだけです.

おまけ

問題の変種を考えます.

変種 1

タイプ 2 のクエリを次のように変更することが考えられます.

  • 2 v d : 頂点 $v$ からの最短距離が $d$ 以下 であるような頂点の集合を $S(v,d)$ として,$\displaystyle \sum _ {u \in S(v,d)} W _ u$ を出力する.

元の問題と同様の方法で解こうとすると,下図のように,集約クエリにおいて $\mathrm{Seg} _ p$ で和を取る必要のある区間 (青で囲った部分) の個数が定数個ではなくなってしまいます.

集約クエリ

そこで,下図のように,一旦必要の無い部分を含めた総和 (青で囲った部分) を求めた後に不要な部分 (赤で囲った部分) を差し引くことを考えます.

集約クエリ (改善版)

$p$ の全ての子 $c$ に対して,$U _ p$ における $c$ を根とする子部分木の頂点を BFS の訪問順に並べ替えて Segment Tree $\mathrm{Seg} _ {p, c}$を構築しておくことで,差し引くべき値は $\mathrm{Seg} _ {p, c}$ におけるある区間和として表現され,元の問題と同じ計算量オーダーで問題を解くことが出来ます.

ただし,元の問題は集約クエリの二項演算が可換モノイドであることのみを要求していましたが,この変種においては加えて可逆性が必要であり,即ち可換群であることを要求します 2

変種 1'

変種 1 の自然な拡張です.タイプ 2 のクエリを次のように変更します.

  • 2 v l r : 頂点 $v$ からの最短距離が $l$ 以上 $r$ 以下であるような頂点の集合を $S(v,l,r)$ として,$\displaystyle \sum _ {u \in S(v,l,r)} W _ u$ を出力する.

可逆性を仮定すれば,($l$ 以上 $r$ 以下の総和) = ($r$ 以下の総和) - ($l - 1$ 以下の総和) として解けます.

変種 2

タイプ 3 のクエリを追加します.Segment Tree -> Lazy Segment Tree みたいな気分です.

  • 3 v d f : 頂点 $v$ からの最短距離が $d$ であるような頂点の集合を $S(v,d)$ として,全ての $u \in S(v,d)$ に対して $W _ u \leftarrow f \cdot W _ u$ とする.

これは,現在の方法では,全ての $\mathrm{Seg} _ v$ に正しい情報が乗った状態を効率的に保つのが難しいです.高速に解けたら教えてください.


  1. この方法でも $O(\sqrt{N})$ 個の区間にはなるので愚直よりは良い計算量オーダーが得られているように思われますが,その $O(\sqrt{N})$ 個の区間を高速に求められるかはまた別の問題で,これは結構難しそうだと思っています.
  2. 可逆性を要求せずに計算量オーダを悪化させずに解く方法は思いつきませんでした.出来たら教えてください.

AtCoder Beginner Contest 162 E - Sum of gcd of Tuples (Hard)

はじめに

$O(K ^ \frac{2}{3} (\log \log K) ^ \frac{1}{3} + K ^ \frac{1}{2} \log N)$ 解法を解説します。

解法

過程を色々すっ飛ばすと、答えは式 $(1)$ で計算されることが分かります。ここで、$\varphi$ は オイラーの $\varphi$ 関数 です。

$$ \sum _ {i = 1} ^ K \varphi(i) \cdot\left\lfloor\dfrac{K}{i}\right\rfloor ^ N \tag{1} $$

この式にたどり着くまでの考察などは以下の記事で解説されていますので、そちらを参照してください。

qiita.com

さて、式 $(1)$ に現れる $\displaystyle \left\lfloor\dfrac{K}{i}\right\rfloor$ について、$i$ を正整数全体で動かしたときに高々 $O(\sqrt{K})$ 種類の値しか取らないことが知られています。

証明

まず、$1\leq i \leq \sqrt{K}$ に対して、ありうる $\displaystyle \left\lfloor\dfrac{K}{i}\right\rfloor$ の値は高々 $\lfloor \sqrt{K}\rfloor$ 種類です。

また、$\sqrt{K}\lt i$ に対しては、$\dfrac{K}{i}\lt \sqrt{K}$ が成り立つのでありうる $\displaystyle \left\lfloor\dfrac{K}{i}\right\rfloor$ の値は高々 $\lfloor \sqrt{K}\rfloor$ 種類です。

以上より、$i$ が正整数全体を動くとき、$\displaystyle \left\lfloor\dfrac{K}{i}\right\rfloor$ が取り得る値の種類は高々 $2\lfloor\sqrt{K}\rfloor=O(\sqrt{K})$ 通りであることが示されました。(証明終)

そこで、式 $(1)$ の和を $\displaystyle \left\lfloor\dfrac{K}{i}\right\rfloor$ の値ごとにまとめて計算することを考えます。$\displaystyle \left\lfloor\dfrac{K}{i}\right\rfloor$ が取り得る値の集合を $S$、正整数 $q\in S$ に対して $\displaystyle \left\lfloor\dfrac{K}{i}\right\rfloor = q$ となる正整数 $i$ の集合を $S _ q$ とします。

$$ \begin{aligned} i \in S _ q &\iff q \leq \dfrac{K}{i} \lt q + 1 \land i \in \N^+ \\ &\iff \dfrac{K}{q+1}\lt i \leq \dfrac{K}{q} \land i \in \N^+ \\ &\iff \left\lfloor\dfrac{K}{q+1}\right\rfloor \lt i \leq \left\lfloor\dfrac{K}{q} \right\rfloor \land i \in \N^+ \end{aligned} $$

より、$\displaystyle f(q) := \sum _ {i = 1} ^ {\lfloor K/q \rfloor} \varphi(i)$ と定義すれば、式 $(1)$ は次の式 $(2)$ のように書き換えられます。

$$ \sum _ {q \in S} (f(q) - f(q + 1)) \cdot q ^ N \tag{2} $$

$f(\cdot)$ が $O(1)$ で取得できると仮定すれば、$|S| = O(K ^ \frac{1}{2})$ より式 $(2)$ は $O(K ^ \frac{1}{2} \log N)$ で計算することが出来ます。

実は、$f(q)$ は以下の記事で解説されているアルゴリズムにより $O(K ^ \frac{2}{3} (\log \log K) ^ \frac{1}{3})$ で全て前計算することが出来ます。

yukicoder.me

以上より、本問を全体 $O(K ^ \frac{2}{3} (\log \log K) ^ \frac{1}{3} + K ^ \frac{1}{2} \log N)$ で解くことが出来ました。

実装

atcoder.jp

おわりに

独自要素が一切無くて、これで良いのかという気はする。

ARC034 D - インフレゲーム

問題

atcoder.jp

解法

期待値ではなく,スコアの総和を考えることにします.スコアの総和 $S$ が求まれば,答えは $\dfrac{ S } { ( A + B + C ) !}$ として容易に計算されます.

まずは具体的な例を考えます.

$a _ 1, a _ 2, b _ 1, b _ 2, a _ 3, b _ 3, c _ 1$ の順にカードを引いたとき,スコアは $( ( a _ 1 + a _ 2 ) b _ 1 b _ 2 + a _ 3 ) b _ 3 = a _ 1 \cdot b _ 1 b _ 2 b _ 3 + a _ 2 \cdot b _ 1 b _ 2 b _ 3 + a _ 3 \cdot b _ 3$ となります.従って,$a _ i$ の係数が $b _ { x _ 1 } \cdots b _ { x _ k } \; (1 \leq x _ 1 \lt \cdots \lt x _ k \leq B)$ となるようなカードの並べ方の数を $f _ i (x _ 1, \ldots, x _ k)$ とおけば,スコアの総和 $S$ は次で求められます.

$$ S = \sum _ { i = 1 } ^ A \sum _ { 1\leq x _ 1 \lt \cdots \lt x _ k \leq B } a _ i \cdot f _ i (x _ 1, \ldots, x _ k) \cdot \prod _ {j = 1} ^ k b _ { x _ j } $$

$f _ i (x _ 1, \ldots, x _ k)$ の具体的な表示について考えます.

まず,山札において,$a _ i$ の後に $b _ { x _ 1 }, \ldots, b _ { x _ k }$ が順不同で並んでおり,その後に $c _ 1, \ldots, c _ C$ が順不同で並んでいる必要があります.

ここで,$a _ i$ の並べ方は $1$ 通り,$b _ { x _ 1 }, \ldots, b _ { x _ k }$ の並べ方は $k !$ 通り,$c _ 1, \ldots, c _ C$ の並べ方は $C !$ 通りです.

$b _ { x _ 1 }, \ldots, b _ { x _ k }$ や $c _ 1, \ldots, c _ C$ はどのように並んでいても等価なので,列 $(a _ i, b _ { x _ 1 }, \ldots, b _ { x _ k }, c _ 1, \ldots, c _ C)$ に対して残りの $a$ や $b$ のカードを挿入する方法の数を考えます.

まだ列に含まれていない $b$ のカードは $B - k$ 枚あり,これらは $a _ i$ よりも左か,または $c _ 1$ よりも右に並べられている必要があります.従って,$a _ i, b _ {x _ 1}, \ldots, b _ {x _ k}, c _ 1$ をまとめて $1$ 枚のカードとして見ることで,まだ列に含まれていない $B - k$ 枚の $b$ のカードを挿入する方法は $\displaystyle \binom{B - k + C}{C} \cdot (B - k) ! = \dfrac{(B - k + C) !}{C !}$ 通りあることが分かります.

残りの $A - 1$ 枚の $a$ のカードは $a _ i$ の係数には影響しないのでどこに挿入してもよく,即ち $\displaystyle \binom{A + B + C}{A - 1} \cdot (A - 1) ! = \dfrac{(A + B + C) !}{(B + C + 1) !}$ 通りの挿入のしかたがあります.

以上より,次の $f _ i (x _ 1, \ldots, x _ k)$ の表示を得ます.重要な事実として,$i$ に依存しないことを注意しておきます.

$$ \begin{aligned} f _ i (x _ 1, \ldots, x _ k) &= k! \cdot C! \cdot \dfrac{(B - k + C) !}{C !} \cdot \dfrac{(A + B + C) !}{(B + C + 1) !} \\ &= \dfrac{(A + B + C) !}{B + C + 1} \cdot \binom{B + C}{k} ^ {-1} \end{aligned} $$

従って,答えは次のように計算されます.

$$ \begin{aligned} \dfrac{S}{(A+B+C)!}&=\dfrac{1}{B + C + 1} \sum _ { i = 1 } ^ A \sum _ { 1\leq x _ 1 \lt \cdots \lt x _ k \leq B } a _ i \cdot \binom{B + C}{k} ^ {-1} \cdot \prod _ {j = 1} ^ k b _ { x _ j } \\ &= \dfrac{1}{B + C + 1} \sum _ {i = 1} ^ A a _ i \sum _ {k = 0} ^ B \binom{B + C}{k} ^ {-1} \cdot T _ k \end{aligned} $$

途中で $\displaystyle T _ k = \sum _ { 1\leq x _ 1 \lt \cdots \lt x _ k \leq B }\prod _ {j = 1} ^ k b _ { x _ j }$ としました.$T _ k$ は,$\displaystyle F(x) := \prod _ {i = 1} ^ B (1 + b _ i x)$ の $x ^ k$ の係数として計算されます.$F(x)$ はナイーブな DP で $\Theta(B ^ 2)$ で,あるいは FFT を用いることで $\Theta(B (\log B) ^ 2)$ で計算できます.

$\displaystyle \sum _ {i = 1} ^ A a _ i$ の計算は $\Theta(A)$ で,式中に現れる二項係数は $\Theta(B)$ で前計算可能なので,本問を $\Theta(A + B ^ 2)$ や $\Theta(A + B (\log B) ^ 2)$ で解くことが出来ました.

感想

現代なら $\mathrm{mod}\; 998244353$ 出力で $1 \leq A, B, C\leq 2\cdot 10 ^ 5, 0 \leq a _ i, b _ i \lt 998244353$ みたいな制約になりそう.$C$ は $B + C + 1 \lt 998244353$ の範囲で大きくできるけど,制約が不自然になるから難しそう (やるなら $A + B + C \lt 998244353$ とか?).