Codeforces Global Round 6 F. Almost Same Distance

問題概要

$n$ 頂点の木が与えられる。いくつかの頂点からなる集合 $W$ が正整数 $k$ に対して以下の条件を満たすとき、 $W$ は almost $k$ -uniformであると呼ぶことにする。

$W$ に含まれる異なる2頂点のペア全てにおいて、その距離は $k$ または $k+1$ である。

$k=1, ..., n$ それぞれに対して、almost $k$ -uniformである集合の最大サイズを求めよ。

制約

$2 \le n \le 5\times 10^{5}$
与えられるグラフは木である

解法

集合 $W$ がalmost $k$ -uniformであることを、簡単のため「 $W$ は条件 $U(k)$ を満たす」と書くことにします。頂点 $a$ の次数を $\deg(a)$ と表記します。頂点 $a, b$ 間の距離を $D(a, b)$ と表記します。

特別ケースを処理

まずはいくつかの特別ケースを処理しておきます。

1頂点からなる集合は「異なる2頂点のペア」が存在しないため、自動的に任意の $k$ に対して $U(k)$ を満たします。

$k=1$ の場合を特別に考えます。「ある1頂点とその隣接頂点全てからなる集合」およびその部分集合は必ず $U(1)$ を満たし、また $U(1)$ を満たす集合はこのようなものに限られます。そのため $k=1$ に対する答えは、頂点の次数の最大値に1を足したものになります。

以降は $k \ge 2$ に対して考え、2頂点以上からなる集合だけを考慮します。

頂点を基準に考えて作る集合

例えば「ある1つの頂点 $r$ から見て、経路が重複せず等距離にあるような頂点の集合」を考えます。この距離が $d$ であるとき、集合内の異なる2頂点間の距離は全て $2d$ になります。

f:id:betrue12:20191218231931p:plain

このように、「ある頂点 $r$ を根と見なし、 $r$ のそれぞれの子が持つ部分木から高々1個ずつの頂点を選ぶ」という方法で作る集合について考えます。

実際には距離が $k$ または $k+1$ であれば $U(k)$ を満たすので、 $r$ からの距離が全て等しい必要はありません。具体的には以下の3パターンが考えられます。

$r$ からの距離が $\lbrace d-1, d, ..., d \rbrace$ となるように頂点を選ぶ。このとき2頂点間の距離は $2d-1$ または $2d$ なので、集合は $U(2d-1)$ を満たす。
$r$ からの距離が全て $d$ となるように頂点を選ぶ。このとき2頂点間の距離は全て $2d$ なので、集合は $U(2d)$ を満たす。
$r$ からの距離が $\lbrace d, ..., d, d+1\rbrace$ となるように頂点を選ぶ。このとき2頂点間の距離は $2d$ または $2d+1$ なので、集合は $U(2d)$ を満たす。

このうちパターン3については考える必要がありません。 $d+1$ の頂点を1つ浅い頂点に置き換えることで、必ずパターン2に変換できるからです。

またこのようにして作った集合が $U(k)$ を満たすならば、頂点を適切に浅いものと置き換えることで、必ず同サイズで $U(k-1)$ を満たす集合を作ることができます。そのため使う頂点の個数を決めたら、なるべく大きい $k$ について $U(k)$ を満たすように頂点を選べば良いです。

ということで、 $2 \le s \le \deg(r)$ であるそれぞれの $s$ について、 $s$ 個の部分木から頂点を1つずつ選んでサイズ $s$ の集合を作ることを考えます。この集合が $U(k)$ を満たすような最大の $k$ は以下のように求めることができます。

$r$ の子が持つ部分木それぞれについて、その部分木内で最も $r$ から遠い頂点との距離を計算する。
上記の距離を大きい方から $s$ 個取り、その最小値を $d_{\min}$ とする。
$s$ 個の距離の中に $d_{\min}$ が1個だけ含まれる場合、最大で $U(2d_{\min} + 1)$ を満たすような集合が取れる。そうでない場合、最大で $U(2d_{\min})$ を満たすような集合が取れる。

これは木DPで計算できます。各頂点 $v$ について「 $v$ の子が持つ部分木それぞれについての、その部分木内で最も $v$ から遠い頂点との距離を集めたもの」を $dp\lbrack v \rbrack$ とし、その最大値+1を順次親に渡していきます。根 $r$ に情報が集まったら、 $dp\lbrack r \rbrack$ を用いて上記の計算ができます。実際には $r$ を全通り試す必要がありますが、これは全方位木DPで処理することができます。

この集合の作り方で、各 $k$ について求めた「 $U(k)$ を満たす集合の最大サイズ」を $ans_{1}\lbrack k \rbrack$ とします。先ほど書いたように集合が $U(k)$ を満たすならば必ず同サイズで $U(k-1)$ を満たせるので、 $r$ を全通り処理した後は $k$ が大きい順に「 $ans_{1}\lbrack k \rbrack$ を $\max(ans_{1}\lbrack k \rbrack, ans_{1}\lbrack k+1 \rbrack)$ に置き換える」という処理をしてやれば良いです。

計算量については、 $\deg(r)$ の総和が $O(n)$ なので、各 $r$ に対して $s$ を1つずつ処理したとしても合計で $O(n)$ 通り。距離のソート等を考慮しても $O(n\log n)$ でできるので計算量は大丈夫です…今のところは。

辺を基準に考えて作る集合

実はこれまで考えたパターンだけでは考慮できないケースが存在します。以下の図のような場合です。

f:id:betrue12:20191219000249p:plain

これは「ある1本の辺を基準として、その両端の頂点 $r, c$ が持つ部分木から、経路が重複せずに $r, c$ それぞれからの距離が全て等しく $d$ である頂点を選んだもの」です。この集合の2点間距離は $2d$ または $2d+1$ になり、 $U(2d)$ を満たします。これは先ほどのパターンでは考慮できないものです。

これも同じように全方位木DPで処理します。根 $r$ とその子 $c$ を繋ぐ辺を見る時には、先ほどの $dp\lbrack r \rbrack$ （から $c$ 以下の部分木に由来する値を除いたもの）と $dp\lbrack c \rbrack$ を使います。これらの距離を混ぜて大きい方から $s$ 個取り、その最小値を $d_{\min}$ とすれば、 $U(2d_{\min})$ を満たすサイズ $s$ の集合を構成できます。

またこのようにして作った $U(2d)$ を満たす集合からは、全ての頂点を1つ浅いものに入れ替えることで $U(2d-2)$ を満たす集合を作ることができます。 $U(2d-1)$ を満たす集合にはできません。

これで先ほどと同じように解けたように見えますが、問題は計算量です。各 $r, c$ のペアについてサイズ $s$ を1つずつ試すと、次数について2乗オーダーの計算量が掛かるので間に合いません。

計算量削減のため、 $dp\lbrack v \rbrack$ として個々の値を単純に並べるのではなく、「値○が△個」という形にした頻度表で持つことを考えます。 $dp\lbrack r \rbrack$ と $dp\lbrack c \rbrack$ を混ぜた時、そこに含まれる距離を合計した値は $O(n)$ です。つまりその種類数は $O(\sqrt{n})$ になります。

集合のサイズとして $s$ の値を全て試す必要はなく、以下のように各種類ごとにまとめて処理することで、1本の辺についての処理が $O(\sqrt{n})$ で完了し、全体 $O(n\sqrt{n})$ で計算することができます。

f:id:betrue12:20191219005020p:plain

この集合の作り方で、各 $k$ について求めた「 $U(k)$ を満たす集合の最大サイズ」を $ans_{2}\lbrack k \rbrack$ とします。先ほど書いたように集合が $U(k)$ を満たすならば必ず同サイズで $U(k-2)$ を満たせるので、 $r$ を全通り処理した後は $k$ が大きい順に「 $ans_{2}\lbrack k \rbrack$ を $\max(ans_{2}\lbrack k \rbrack, ans_{2}\lbrack k+2 \rbrack)$ に置き換える」という処理をしてやれば良いです。

まとめ

こうして「頂点を基準に考えて作る集合」と「辺を基準に考えて作る集合」、そして冒頭に書いたいくつかの特別ケースを考えてきました。これで考えるべきパターンが尽くされているため、 $k$ ごとにそれぞれで求めた集合サイズの最大値を取ることで答えが求められます。

考えるべきパターンがこれで尽くされていることの証明

Editorialには書いていなかったので証明を付けておきます（がんばりました）。

これまでの説明ではパターンを天下り的に列挙してきましたが、考えるべき集合 $W$ がこれらのパターンで尽くされていることを確認する必要があります。「特別ケース」として扱った $k=1$ の場合とサイズ1の集合については良いとして、 $k \ge 2$ に対して $U(k)$ を満たすサイズ2以上の集合 $W$ が、先に説明した「頂点基準」あるいは「辺基準」の作り方で作られる集合のいずれかに合致することを示します。

まず $k \ge 2$ を仮定すると、 $W$ に含まれる3頂点以上が1本のパスに含まれることはありません。この3頂点を並んだ順に $a, b, c$ とすると、 $D(a, c) = D(a, b) + D(b, c) \ge D(a, b) + 2$ となり条件を満たさないからです。