AtCoder Beginner Contest 160 D,E,F問題メモ

AtCoder Beginner Contest 160

D - Line++

• $N$ 頂点のグラフ、辺は以下の $N$ 本
  • 各頂点 $i$ と $i+1$ を結ぶ $N-1$ 本
  • 頂点 $X$ と $Y$ を結ぶ1本
• 最短距離が $k$ である頂点の組の個数を、各 $1 \le k \le N-1$ について求めよ
• $1 \le N \le 2 \times 10^3$
• $X+1 \lt Y$

こんな感じのグラフ。

1 -- 2 -- 3 -- 4 -- 5 -- 6 -- 7 -- 8
      `-----------------'

制約が小さく $O(N^2)$ で大丈夫なので、全ての組の最短距離を調べればよい。

$i$ から $j$（$i \lt j$）に行く移動パターンは、以下の2つしかない。

• $i→j$（短縮辺を使わない）
• $i→X→Y→j$

これは、それぞれ以下で計算できる。

• $j-i$
• $|i-X|+1+|j-Y|$

よって、この2つの小さい方が $(i,j)$ の最短距離なので、最短距離毎に数え挙げればよい。

E - Red and Green Apples

• 赤いリンゴが $A$ 個、緑のリンゴが $B$ 個、無色のリンゴが $C$ 個ある
• それぞれ美味しさが $(p_1,...,p_A),(q_1,...,q_B),(r_1,...,r_C)$ で与えられる
• ここから赤いリンゴを $X$ 個、緑のリンゴを $Y$ 個食べる
• 無色のリンゴは、赤いリンゴとしても緑のリンゴとしても食べられる
• 食べるリンゴの美味しさの和の最大値を求めよ
• $1 \le A,B,C \le 10^5$

Eにしてはすごく単純なのに、パッと見えなかったのはよくない。

まず、美味しい方から並べて $X+1$ 番目以下の赤いリンゴと $Y+1$ 番目以下のリンゴは絶対に採用されないので除く。

後は無色のリンゴの方が美味しければ、赤や緑のリンゴの美味しくない方から順に置き換えていけばよい。

これは単に、上記で残った赤緑のリンゴと無色のリンゴを全て合わせ、美味しい方から $X+Y$ 個取ればよい。

F - Distributing Integers

• $N$ 頂点の木、$i$ 番目の辺は $a_i$ と $b_i$ を結ぶ
• 全ての頂点 $k$ について、以下を求めよ
  • 頂点 $k$ に $1$ を書き込む
  • 残りの頂点に $2,3,...,N$ を順に書き込んでいく
  • この時、既に数字が書き込まれた頂点に隣接する頂点にしか書き込めない
  • 書き込み方のパターン数を $\mod 10^9+7$ で求めよ
• $2 \le N \le 2 \times 10^5$

Eの単純さから比べると急激な実装難度の差。 ただ、Typical DP Contestなどで類題があったらしい？ので、流用できる部分は流用できたか。

こういう「全ての頂点について、そこを根としたときの何か」を求めるときには、全方位木DPを用いる解法がある。

• まず頂点1を根として、DFSなどで、各頂点 $v$ について $v$ 以下の部分木のみを考えた答えを求める
• 再度DFSなどで親方面の部分木の情報を与えながら遷移を行い、適切にマージすることで、$v$ を木全体の根とした時の答えが求まる
  • 親方面の部分木の情報は、1回目の計算結果からある程度高速に計算できる

実装は軽くは無いが、割とテンプレ化しやすく、何をすればよいかはわかりやすい。 （典型的な問題でABC-Fに置かれるくらいの難度はあるので、それを更に発展した問題は出会う機会が少ないというだけかも）

• 【全方位木DP】明日使える便利な木構造のアルゴリズム - Qiita

以下、「頂点番号」と「そこに書き込む数字」の2種類を区別するため、頂点番号を①、書き込む数字を❶で表す。

求めるのは書き込むパターン数だが、具体的に書き込む数字は無視して、順番だけ考えれば書き込む数字に1対1対応できる。

①--②--③--④
     |   `--⑤--⑥
     `--⑦--⑧

①を根とした時の、③以下に書き込むパターン数は
❶--❷            ❶--❸            ❶--❹
 `--❸--❹         `--❷--❹         `--❷--❸

書き込む頂点の順番で表現する
③→④→⑤→⑥    ③→⑤→④→⑥    ③→⑤→⑥→④

同様に、⑦以下に書き込むパターン数は以下の1通り
⑦→⑧

さて、ここで2つの部分木をマージする方法について考える。 以下の2段階に分けて考えることができる。

• (A) まず、何番目にはどちらの部分木に書き込むか、のパターン数を求める
  • 一方の頂点の並びに、もう一方を挟み込む感じ
• (B) それぞれの部分木の中での書き込み方のパターン数を掛け合わせる

③を根とする部分木の書き込み方  ○→○→○→○  パターン数 3
⑦を根とする部分木の書き込み方  ○→○          パターン数 1

4個のボールの並び（両端含む5箇所）に、2個のボールを挟み込む場合の数
↓  ↓  ↓  ↓  ↓
  ○  ○  ○  ○
重複組み合わせ 5H2 = (5+2-1)C(2)  ... これが(A)

それぞれの具体的な頂点パターン:
  2つの部分木の書き込み方の積 3 x 1  ... これが(B)

これらをまとめて、頂点②以下の書き込み方のパターン数
5H2 x 3 x 1 = 15 x 3 x 1 = 45

これは、子が3個以上あっても順にマージしていけばよい。

マージに必要な情報を考えると、$v$ 以下の「頂点数」と「書き込み方のパターン数」があればよいと分かる。

葉の頂点数は0、パターン数は1なので、ここからボトムアップに求めていけばよい。 （計算上、頂点数は $v$ 自身を含めない方が何かと便利なので、そうしている）

①--②--③--④
     |   `--⑤--⑥
     `--⑦--⑧

②から③への遷移について考える。

③を根とした木のイメージ

③--④          ┬→(A) 1回目で計算済み
 |--⑤--⑥      ┘
 `--②--①      ┬→(B) 親（②）から与えたい
     `--⑦--⑧  ┘

(A)と(B)をマージすれば、③を根としたときの答えが求まる

(B)の作り方:

    ②--③--④          ②の部分木として1回目で計算済み
     |   `--⑤--⑥
     `--⑦--⑧

    ②                  ③の部分木を除外する
     `--⑦--⑧
    
    ②--①              ②の親（①）方面の部分木をマージする
     `--⑦--⑧

(B)を作るために、以下をマージしたい。

• ②にとっての親方面の部分木（②を根と考えたときの①以下の部分木）
• ②の子（③、⑦）の内、遷移先（③）を除いた全ての部分木のマージ結果

1つめは親から与えられるため、所与としてよい。 2つめが厄介で、上手いこと③の部分木だけ除いた結果を再計算したい。 （毎回1つずつマージしていては、たとえばスター_(グラフ理論)のような木の場合、ほぼ $N$ 頂点のマージを $N$ 回行うことになり、TLE）

頂点 $v$ 以下の頂点数を $C_v$、書き込み方のパターン数を $P_v$ で表すと、マージを逆算して

$P_② = {}_{C_② - C_③ - 1}H_{C_③} \times P_③ \times P_{③以外}$ より、

$P_{③以外} = \dfrac{P_②}{{}_{C_② - C_③ - 1}H_{C_③} \times P_③}$ となる。

ただ今回、MODを取っているので割り算ができない。

こういう場合は、前後からの累積マージ結果を持っておく方法がある。

                                      ③  ⑦
前からマージした場合のカウント    [0,  4,  6    ]
前からマージした場合のパターン数  [1,  3, 45    ]
後からマージした場合のカウント    [    6,  2,  0]
後からマージした場合のパターン数  [   45,  1,  1]

③を除いたマージ結果を求めたいときは、以下の2つを1回だけマージすればよい。

• 前からの③より手前のカウントとパターン数 (0, 1)
• 後からの③より後ろのカウントとパターン数 (2, 1)

ここに、親からの情報をマージすれば、頂点③に渡す親方面の情報となる。

（これを実現するためには、あらかじめ辺リストから親頂点を除き、子の順番を固定しておく）

ただ、③以外の頂点のマージ結果を逆算する箇所で、MODを取っているので割り算はできないと言ったが、 フェルマーの小定理より(MOD-2)乗すれば逆数となるので、 下手に実装を複雑化するより、毎回(MOD-2)乗を計算した方が速かったかも。