こんな感じのグラフ。
1 -- 2 -- 3 -- 4 -- 5 -- 6 -- 7 -- 8
`-----------------'
制約が小さく $O(N^2)$ で大丈夫なので、全ての組の最短距離を調べればよい。
$i$ から $j$（$i \lt j$）に行く移動パターンは、以下の2つしかない。
これは、それぞれ以下で計算できる。
よって、この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にしてはすごく単純なのに、パッと見えなかったのはよくない。
まず、美味しい方から並べて $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]))
Eの単純さから比べると急激な実装難度の差。 ただ、Typical DP Contestなどで類題があったらしい？ので、流用できる部分は流用できたか。
こういう「全ての頂点について、そこを根としたときの何か」を求めるときには、全方位木DPを用いる解法がある。
実装は軽くは無いが、割とテンプレ化しやすく、何をすればよいかはわかりやすい。 （典型的な問題でABC-Fに置かれるくらいの難度はあるので、それを更に発展した問題は出会う機会が少ないというだけかも）
以下、「頂点番号」と「そこに書き込む数字」の2種類を区別するため、頂点番号を①、書き込む数字を❶で表す。
求めるのは書き込むパターン数だが、具体的に書き込む数字は無視して、順番だけ考えれば書き込む数字に1対1対応できる。
①--②--③--④
| `--⑤--⑥
`--⑦--⑧
①を根とした時の、③以下に書き込むパターン数は
❶--❷ ❶--❸ ❶--❹
`--❸--❹ `--❷--❹ `--❷--❸
書き込む頂点の順番で表現する
③→④→⑤→⑥ ③→⑤→④→⑥ ③→⑤→⑥→④
同様に、⑦以下に書き込むパターン数は以下の1通り
⑦→⑧
さて、ここで2つの部分木をマージする方法について考える。 以下の2段階に分けて考えることができる。
③を根とする部分木の書き込み方 ○→○→○→○ パターン数 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回だけマージすればよい。
ここに、親からの情報をマージすれば、頂点③に渡す親方面の情報となる。
（これを実現するためには、あらかじめ辺リストから親頂点を除き、子の順番を固定しておく）
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)乗を計算した方が速かったかも。