AtCoder Beginner Contest 133 E - MEMORANDUM CEDRETABRI

E - Virus Tree 2

この問題はコンテスト中には時間が足りなくて解けなかったのだけど、その後解いたので、何を考えたかを書いておきたい。

こういう木の問題、僕は割と好きだ。

概要

木がある。各頂点を ${K}$ 色以下に塗り分けたい。ただし、頂点同士の距離が ${2}$ 以下の場合は別の色にしなければならない。さて、何通りの塗り分け方が考えられるだろうか。

よくある ${\mod10^9+7}$ で組み合わせの数を求める問題だ。

何を考えたか

ここでは、コンテスト中（と、終わった後少し）に僕が考えたことをつらつらと書いていく。

当時、ここまで筋道だって考えたわけではないと思うが。

順列を考える

まず、組み合わせについて色々考えてみた。

上のような条件無しで、同じ色は ${1}$ 度しか使えないこととして、 ${N}$ 個の頂点を ${K}$ 色で塗り分ける方法はいくつあるだろうか？

これは順列なので、 ${{}_K\mathrm{P}_N}$ で求めることができるのはわかる。

これをもう少し解きほぐしていくと、

頂点の色を選ぶ
- ${K}$ 色残っているので、 ${K}$ 通りの選び方がある
頂点の色を選ぶ
- ${K-1}$ 色残っているので、 ${K-1}$ 通りの選び方がある
頂点の色を選ぶ
- ${K-2}$ 色残っているので、 ${K-2}$ 通りの選び方がある
...
頂点の色を選ぶ
- ${K-N+1}$ 色残っているので、 ${K-N+1}$ 通りの選び方がある

それで結局、 ${\dfrac{K!}{(K-N)!}}$ というご存知の公式が出てくるわけだが、ここでわかることは、

ある頂点で色が決まると、他の頂点が使える色の数は ${1}$ つ少なくなる
全体の場合の数は、各頂点が選べる数の総乗

ということである。

（ここまでは多分中学生の数学の内容だと思う。のだが、僕の記憶は怪しかった……。）

距離 ${2}$ 以下の頂点

ここで、問題を振り返ってみる。

上の条件だと、「その色を使えるのは全ての頂点で ${1}$ 回だけ」としたので ${1}$ つの頂点で色を決めたら他の全ての頂点で使える色が ${1}$ 少なくなったわけだが、今回の問題の条件だとそこが少し違っていて、

距離が ${2}$ 以下の頂点で同じ色が使えない

これはつまり、

ある頂点で色を決めると、距離が ${2}$ 以下の頂点で使える色が ${1}$ つ少なくなる

と言い換えることができる。

そして、「全体の場合の数は、各頂点が選べる数の総乗」という点に変化はないはずなので、上手いこと各頂点の「使える色」を決めていってやればこの問題は解けそうである。

しかし、距離 ${2}$ 以下の使える色を制限する、という状態をどのように持たせればいいだろうか？

最初は、全ての頂点に仮に ${K}$ の色を持たせておき、各頂点で色を決めたら、そこから距離 ${2}$ 以下の頂点全ての色の数を減らす、みたいなことを考えたが、これだと計算量が大きくなりすぎる。

次に、各頂点の色を決める時に、距離 ${2}$ までの頂点を見に行って、色が決まっていたら自分が使える色を ${1}$ つ減らす、みたいなことも考えたが、これもやはり計算量が大きくなりすぎる。

というわけで「配る」アプローチと「貰う」アプローチとの両方を考えたが、この ${2}$ つは結局やっていることは同じなので、どちらも上手くいかない。

なのだが、その両方を組み合わせるとどうだろう？

中間地点

ある頂点で色を決めることを考える。もし距離 ${2}$ 以下の頂点のどれもがまだ色を決めていなければ、この頂点は ${K}$ 色の中から色を決めることができる。

そして色を決めた後で、自分に隣接する頂点の「隣接した頂点で色が決まっている数」カウンタに ${1}$ 加える。こいつが管理する数は、名前の通り隣接する頂点で決まった色の数だけだ（距離 ${2}$ 以下の頂点全てではない）。

では最初に戻って、ある頂点が色を決める時、距離 ${2}$ 以下の頂点でどれだけ色が決まっているか、どうやって判断すればいいだろうか？

まず、自分自身のカウンタを見る。これは、自分に隣接する頂点で色が決まった数なので、当然距離 ${2}$ 以下に含まれる。

次に、自分に隣接する頂点全てのカウンタを見に行く。これは、自分に隣接する頂点に隣接する頂点、すなわち距離 ${2}$ の頂点で色が決まった数を表す。

この時、自分自身をカウントしてしまう恐れはない。まだ自分の色は決まっていないからだ。また、同じ頂点を重複してカウントしてしまう恐れもない。何故ならグラフは木なので、ある頂点からある頂点への経路は高々 ${1}$ つしか存在しないはずだからだ。

これらカウンタでカウントされている数の合計を ${K}$ から引けば、その頂点で使える色の数が出てくる。

要するに、「色を使った」という情報を中間地点にある頂点を経由して伝播させるわけだ。

この作業の計算量を見積もってみよう。全ての頂点を巡るので ${N}$ 回の処理。かつ、全ての辺を ${2}$ 回巡るので（配る時と貰う時） ${(N-1)\times 2}$ 回の処理。この ${2}$ つはそれぞれ独立しているので、計算量は ${O(N + (N-1)\times 2)}$ つまり ${O(N)}$ で収まりそうだ。

これはなんとかなりそう。

順序の問題

上の説明だけだと、どの頂点をどういう順番で処理していっても問題ないように思える。残念ながらそうではない。根付き木として、根から順に処理していく必要がある。

どうしてそうなるかと言うと、あまり上手い説明はできないのだが、たとえば直径 ${4}$ 以上のある木の色を両側から塗っていったとして（当然、両端のノードで選べる色の数は ${K}$ だ）、それらが重なる付近で矛盾が起きる組み合わせが出てくる。

もう少し考えてみると、ある木の任意の ${2}$ つの頂点は、距離 ${2}$ 以下の頂点同士を通って辿り着けることは自明なので、最初に色を決める頂点以外が ${K}$ 色から選べるということは絶対に起こらないはずなのだ。

そういう矛盾を防ぐためには、根付き木の根から順番に処理していく他ない。

（ごめんなさい、この節は僕の理解力不足でかなりいい加減なことが書いてあります。お叱りとか受けますので、もし変なことが書いてあったら容赦なく突っ込んでください。）

コード

解法が決まればコードに落とし込むだけだ。よくある DFS なので迷うことはなかった。

import std.stdio, std.algorithm, std.conv, std.array, std.string, std.math, std.typecons, std.numeric;

enum P = 10L^^9+7;

long[][10^^5] T;
long[10^^5] D;
bool[10^^5] M;

void main()
{
    auto nk = readln.split.to!(long[]);
    auto N = nk[0];
    auto K = nk[1];
    foreach (_; 0..N-1) {
        auto ab = readln.split.to!(long[]);
        auto a = ab[0]-1;
        auto b = ab[1]-1;
        T[a] ~= b;
        T[b] ~= a;
    }

    long r = 1;
    long[] ss = [0];
    M[0] = true;
    while (!ss.empty) {
        auto i = ss[0];
        ss = ss[1..$];
        auto ir = K - D[i];
        foreach (j; T[i]) {
            if (!M[j]) {
                ss ~= j;
                M[j] = true;
            }
            ir -= D[j];
            ++D[j];
        }
        r = (r * ir) % P;
    }
    writeln(r);
}

感想

好きなタイプの問題だった。

それだけに、コンテスト中に解ききれなかったことの後悔は大きい。

おのれ C 問題……。

概要

何を考えたか

順列を考える

距離 以下の頂点

中間地点

順序の問題

コード

感想

距離 ${2}$ 以下の頂点