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

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)$ の最短距離なので、最短距離毎に数え挙げればよい。

n, x, y = list(map(int, input().split()))

ans = [0] * n
for i in range(1, n):
    for j in range(i + 1, n + 1):
        ans[min(j - i, abs(i - x) + abs(j - y) + 1)] += 1
print('\n'.join(map(str, ans[1:])))

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$ 個取ればよい。

x, y, a, b, c = list(map(int, input().split()))
ppp = list(map(int, input().split()))
qqq = list(map(int, input().split()))
rrr = list(map(int, input().split()))
ppp.sort(reverse=True)
qqq.sort(reverse=True)
ppqq = ppp[:x] + qqq[:y]
ppqq.extend(rrr)
ppqq.sort(reverse=True)
print(sum(ppqq[: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

1回目のDFS

以下、「頂点番号」と「そこに書き込む数字」の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$ 自身を含めない方が何かと便利なので、そうしている）

2回目のDFS

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

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

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

③--④          ┬→(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)

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

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

import sys

sys.setrecursionlimit(10 ** 5 + 5)


# 2乗の木DP

def prepare(n, MOD):
    f = 1
    factorials = [1]  # 0!の分
    for m in range(1, n + 1):
        f *= m
        f %= MOD
        factorials.append(f)
    inv = pow(f, MOD - 2, MOD)
    invs = [1] * (n + 1)
    invs[n] = inv
    for m in range(n, 1, -1):
        inv *= m
        inv %= MOD
        invs[m - 1] = inv

    return factorials, invs


def dfs0(links):
    q = [(0, -1)]
    while q:
        v, p = q.pop()
        links[v].discard(p)
        links[v] = list(links[v])
        q.extend((u, v) for u in links[v])


def dfs1(links, st_count_fwd, st_count_bwd, st_pattern_fwd, st_pattern_bwd):
    q = [(0, -1, 0)]
    while q:
        v, p, o = q.pop()
        if o == 0:
            # 初回訪問
            q.append((v, p, 1))
            q.extend((u, v, 0) for u in links[v])
        else:
            # 子の探索が終わった後の砲門
            curr_count = [0]
            curr_pattern = [1]
            for u in links[v]:
                stc_u = st_count_fwd[u][-1] + 1
                new_pattern = factorials[curr_count[-1] + stc_u] * invs[curr_count[-1]] * invs[stc_u] % MOD
                curr_pattern.append(new_pattern * curr_pattern[-1] * st_pattern_fwd[u][-1] % MOD)
                curr_count.append(curr_count[-1] + stc_u)
            st_count_fwd[v] = curr_count
            st_pattern_fwd[v] = curr_pattern
            curr_count = [0]
            curr_pattern = [1]
            for u in reversed(links[v]):
                stc_u = st_count_fwd[u][-1] + 1
                new_pattern = factorials[curr_count[-1] + stc_u] * invs[curr_count[-1]] * invs[stc_u] % MOD
                curr_pattern.append(new_pattern * curr_pattern[-1] * st_pattern_fwd[u][-1] % MOD)
                curr_count.append(curr_count[-1] + stc_u)
            st_count_bwd[v] = curr_count[::-1]
            st_pattern_bwd[v] = curr_pattern[::-1]


def dfs2(links, st_count_fwd, st_count_bwd, st_pattern_fwd, st_pattern_bwd, ans):
    def sub(v, pc, pp):
        # print(v, pc, pp)
        if pc == 0:
            ans[v] = st_pattern_fwd[v][-1]
        else:
            cc = st_count_fwd[v][-1]
            cp = st_pattern_fwd[v][-1]
            pattern = factorials[cc + pc] * invs[cc] * invs[pc] % MOD
            pattern = pattern * cp * pp % MOD
            ans[v] = pattern

        for i, u in enumerate(links[v]):
            fc = st_count_fwd[v][i]
            fp = st_pattern_fwd[v][i]
            bc = st_count_bwd[v][i + 1]
            bp = st_pattern_bwd[v][i + 1]
            p1 = factorials[fc + bc] * invs[fc] * invs[bc] % MOD
            p1 = p1 * fp * bp % MOD
            p2 = factorials[fc + bc + pc] * invs[fc + bc] * invs[pc] % MOD
            p2 = p2 * p1 * pp % MOD
            # print(v, u, fc, fp, bc, bp, p1, p2)
            sub(u, fc + bc + pc + 1, p2)

    sub(0, 0, 1)


n, *ab = map(int, sys.stdin.buffer.read().split())
links = [set() for _ in range(n)]
for a, b in zip(ab[0::2], ab[1::2]):
    a -= 1
    b -= 1
    links[a].add(b)
    links[b].add(a)
MOD = 10 ** 9 + 7
factorials, invs = prepare(n, MOD)
st_count_fwd = [[] for _ in range(n)]
st_count_bwd = [[] for _ in range(n)]
st_pattern_fwd = [[] for _ in range(n)]
st_pattern_bwd = [[] for _ in range(n)]
ans = [0] * n

dfs0(links)
dfs1(links, st_count_fwd, st_count_bwd, st_pattern_fwd, st_pattern_bwd)
# print(st_count_fwd, st_count_bwd, st_pattern_fwd, st_pattern_bwd, sep='\n')
dfs2(links, st_count_fwd, st_count_bwd, st_pattern_fwd, st_pattern_bwd, ans)
print('\n'.join(map(str, ans)))

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

目次