目次
AtCoder Beginner Contest 133 D,E,F 問題メモ
D - Rain Flows into Dams
問題
- 円形に $N$ 個の山、時計回りに $1,2,...,N$ が並ぶ
- 山の間にはダムがある。山 $i$ と $i+1$ の間にダム $i$ がある（山$N+1$ は山1を指す）
- 山に降った雨は、両脇のダムに均等に流れる
- ある日の各ダムの増量を調べると、$A_1,A_2,...,A_N$ であった
- ここから、その日の各山に降った雨量を求めよ。ただし雨量は非負整数で偶数とする
- $1 \le N \le 10^5$
- $N$ は奇数
- $0 \le A_i \le 10^9$
解法
適当に1つの山の雨量を仮定すれば、連鎖的に全ての山の雨量が求まる。
山1の雨量を仮定して山2,3,…と決めていく。
山$N$ まで決めたときに、山 $N$ と山1の間のダム$N$ で辻褄が合わなくなるはずなので、そこから最初の仮定の値がどれだけズレていたかを求められる。
山1 ○ 0(仮定)
ダム5 ダム1 5 3 5 3
山5 山2 ○ ○ 6 6
ダム4 ダム2 5 8 5 8
山4 ダム3 山3 ○ 7 ○ 4 7 10
山1の雨量を0と仮定して時計回りに決めていくと、山5の雨量は6となる。つまりこのままではダム5には水量3しか流れ込まない。ダム5の増量は5なので、2足りない。
山1の雨量を $k$ 増やすと、山2の雨量は $k$ 減り、山3の雨量は $k$ 増え……山$N$(奇数)の雨量は $k$ 増える。
よって、ダム$N$ の増量が $k$ 足りない場合、山1の雨量を $k$ 増やすとよい。
2
5 3
8 4
5 8
2 7 12
n = int(input())
aaa = list(map(int, input().split()))
ans = [0]
tmp = 0
for a in aaa[:-1]:
tmp = a - tmp
ans.append(tmp * 2)
adj = aaa[-1] - tmp
ans2 = []
for a in ans:
ans2.append(a + adj)
adj *= -1
print(*ans2)
E - Virus Tree 2
問題
- $N$ 頂点の木があり、頂点には番号 $1,2,...,N$ がつけられている
- $i$ 番目の辺は頂点 $a_i,b_i$ をつなぐ
- 以下の条件を満たして、頂点を $K$ 色以下で塗り分ける方法は何通りか、$\mod{10^9+7}$ で求めよ
- 条件: 2つの異なる頂点 $x,y$ の距離が2以下ならば、頂点 $x$ と頂点 $y$ の色は異なる
- $1 \le N,K \le 10^5$
解法
根を適当に1つ決め、根から「色の自由度」を確定させていく。
K=5
１
／｜＼
２ ３ ４
/\
５６
/\
７８
頂点1を根とする。
親とその全ての子は、距離が2以下なので、互いに同じ色になってはいけない。
頂点1の自由度は $K$ そのままの5。頂点2,3,4の自由度はそれぞれ4,3,2となる。つまり、$K-1$ からスタートして徐々に減っていく。
⑤
／｜＼
④ ③ ②
/\
５６
/\
７８
頂点2,3は行き止まり。頂点4から5,6へ探索を広げる。
今度の子（頂点5,6）は、親である頂点4に加えて、親の親である頂点1とも色が被ってはいけない。
よって、頂点5,6の自由度はそれぞれ3,2となる。つまり、$K-2$ からスタートして徐々に減っていく。
⑤
／｜＼
④ ③ ②
/\
③②
/\
７８
頂点7,8の自由度は、同じように頂点4および5と被ってはいけないので、改めて $K-2$ からスタートして3,2となる。
⑤
／｜＼
④ ③ ②
/\
③②
/\
③②
これらを全て掛け合わせた数が答えとなる。
- 根の自由度は $K$
- 根の直接の子の自由度は、$K-1,K-2,...$
- それ以外の頂点の子の自由度は、$K-2,K-3,...$
- 途中で自由度が0になったら、$K$ 色では塗り分けられないので、探索を打ち切ってよい（最後までやって掛け合わせてもどうせ0になるので打ち切らなくてもいいけど）
import sys
sys.setrecursionlimit(10 ** 5 + 5)
def dfs(v, p, f):
ret = f
child_free = k - 1 if f == k else k - 2
for u in links[v]:
if u == p:
continue
pat = dfs(u, v, child_free)
if pat == 0:
ret = 0
break
ret = ret * pat % MOD
child_free -= 1
# print(v, p, ret, len(links[v]))
return ret
n, k = list(map(int, input().split()))
MOD = 10 ** 9 + 7
RK = pow(k, MOD - 2, MOD)
links = [set() for _ in range(n)]
for line in sys.stdin:
a, b = map(int, line.split())
a -= 1
b -= 1
links[a].add(b)
links[b].add(a)
ans = dfs(0, -1, k)
print(ans)
F - Colorful Tree
問題
- $N$ 頂点の木があり、頂点には番号 $1,2,...,N$ がつけられている
- $i$ 番目の辺は頂点 $a_i,b_i$ をつなぐ。その色は $c_i$、長さは $d_i$ である
- クエリが $Q$ 個与えられるので、それぞれに答えよ
- $i$ 番目のクエリでは、$x_i,y_i,u_i,v_i$ が与えられる
- 色 $x_i$ の全ての辺の長さを $y_i$ に変更した時の、2頂点 $u_i,v_i$ 間の距離を求めよ
- 辺の長さの変更は各クエリで独立で、他のクエリの変更は影響しない
- $1 \le N \le 10^5$
- $1 \le Q \le 10^5$
- $1 \le c_i,x_i \le N-1$
- 実行時間制限: 4sec
解法
本当にABCかと思うくらい実装が重い。一方で、木の2頂点間の距離の求め方さえ知っていれば、何をすればよいかは比較的見えやすい。
単純化した問題
まず、辺の長さの変更が無いとして、ただ任意の2頂点の距離をクエリに応じて求めることを考える。
$N$ が小さければ全ての距離を持っておけばよいが、$10^5 \times 10^5$ の行列は持ってられないので、「最小共通祖先」（LCA: Lowest Common Ancestor）を利用する。
これは、任意の木の2頂点から根に向かって辿ったとき、最初に合流する頂点を求めるアルゴリズムである。
根を適当に1つ決め、そこからの全頂点への距離を求めておく。また、LCAを求められるような前計算データ（上記サイトまたは解説放送などを参照）を作っておく。
１
／｜＼
２ ３ ４
/\
５６
/\
７８
距離
０
／｜＼
10 20 30
/\
70 80
/\
130 140
すると、たとえば頂点6と7の距離を求めたいとき、
- 頂点1～6の距離 $d_6=80$
- 頂点1～7の距離 $d_7=130$
- 頂点6と7のLCA: 頂点4
- 頂点1～4の距離 $d_4=30$
6→1→7と移動した距離210の内の、頂点4→1→4の往復分である30*2=60が無駄なので、210-60=150が頂点6,7の距離、というように求められる。
$dist(6,7)=d_6+d_7-2d_{LCA(6,7)}$
元の問題
あるクエリ $x,y,u,v$ に対して上記の方法で距離を求めようと思うと、次のような情報があればよい。$u,v$ のLCAを $w$ として、
- 根～$u$、根～$v$、根～$w$ それぞれの「色 $x$ の長さが $y$ に変更されたときの距離」
これは、以下を計算しておくと、$A+By$ で求められる。
- (前略) それぞれの「色 $x$ の辺以外の長さの合計(A)と、色 $x$ の辺の本数(B)」
この情報があると、クエリに対して以下のように答えを算出できる。
- $LCA(u,v)=w$ を取得
- $u,v,w$ それぞれの $A_u,A_v,A_w$ と $B_u,B_v,B_w$ を取得
- $(A_u+yB_u)+(A_v+yB_v)-2(A_w+yB_w)$ が答え
ただし、2番目に関して、全ての頂点・全ての色の組合せについてA,Bを求めておくのは無理。 頂点が $N$、色が $N-1$ 個あるので $O(N^2)$ となり計算量的にも空間的にも溢れてしまう。
1つのクエリに必要なのはたかだか1色、3頂点に対しての情報のみなので、どの情報を保持しておけばよいかクエリを先読みして調べることで、間に合うようになる。
- 1回DFSをして、LCAを求めるのに必要なデータを作る
- クエリを先読みし、各クエリについてLCA(頂点 $w$)を求めて「色$x$、頂点$u,v,w$ についての距離情報が必要」という情報を得ておく
- 再度木を探索し、距離情報を必要な部分のみ計算する
- 改めてクエリを読み、クエリ毎の答えを求める
という手順で実装すればよい。
import sys
from collections import defaultdict
sys.setrecursionlimit(10 ** 5 + 5)
# 使用させていただいたLCAのライブラリ
# https://tjkendev.github.io/procon-library/python/graph/lca-segment-tree.html
def solve(n, links, queries):
def _query(a, b, data, m0, INF):
yield INF
a += m0
b += m0
while a < b:
if b & 1:
b -= 1
yield data[b - 1]
if a & 1:
yield data[a - 1]
a += 1
a >>= 1
b >>= 1
def query(u, v):
fu = first_appear[u]
fv = first_appear[v]
if fu > fv:
fu, fv = fv, fu
return euler_tour[min(_query(fu, fv + 1, data, m0, INF))[1]]
def dfs(v, p, dep):
first_appear[v] = len(euler_tour)
depth[v] = dep
euler_tour.append(v)
for u, _, _ in links[v]:
if u == p:
continue
dfs(u, v, dep + 1)
euler_tour.append(v)
def dfs2(v, p, tmp_dists, total_dist):
for x, arr in v_queries[v].items():
arr[0] = total_dist - tmp_dists[x][0]
arr[1] = tmp_dists[x][1]
for u, uc, ud in links[v]:
if u == p:
continue
tmp_dists[uc][0] += ud
tmp_dists[uc][1] += 1
dfs2(u, v, tmp_dists, total_dist + ud)
tmp_dists[uc][0] -= ud
tmp_dists[uc][1] -= 1
INF = (n, None)
euler_tour = []
first_appear = [0] * n
depth = [0] * n
dfs(0, -1, 0)
m = 2 * n
m0 = 2 ** (m - 1).bit_length()
data = [INF] * (2 * m0)
for i, v in enumerate(euler_tour):
data[m0 - 1 + i] = (depth[v], i)
for i in range(m0 - 2, -1, -1):
data[i] = min(data[2 * i + 1], data[2 * i + 2])
v_queries = [{} for _ in range(n)] # [頂点v][色x] = [根～vのx以外の辺の長さ合計, xの辺の本数]
lca = []
for i, (x, y, u, v) in enumerate(queries):
w = query(u, v)
lca.append(w)
v_queries[u][x] = [0, 0]
v_queries[v][x] = [0, 0]
if w != 0:
v_queries[w][x] = [0, 0]
tmp_dists = defaultdict(lambda: [0, 0]) # [色x] = [根～現在の頂点までのxの辺の長さ合計, xの辺の本数] (随時更新)
dfs2(0, -1, tmp_dists, 0)
ans = []
for w, (x, y, u, v) in zip(lca, queries):
qu = v_queries[u][x]
qv = v_queries[v][x]
du = qu[0] + qu[1] * y
dv = qv[0] + qv[1] * y
if w == 0:
ans.append(du + dv)
continue
qw = v_queries[w][x]
dw = qw[0] + qw[1] * y
ans.append(du + dv - 2 * dw)
return ans
n, q = list(map(int, input().split()))
links = [set() for _ in range(n)]
lines = sys.stdin.readlines()
for line in lines[:n - 1]:
a, b, c, d = list(map(int, line.split()))
a -= 1
b -= 1
links[a].add((b, c, d))
links[b].add((a, c, d))
queries = []
for line in lines[n - 1:]:
x, y, u, v = map(int, line.split())
u -= 1
v -= 1
queries.append((x, y, u, v))
ans = solve(n, links, queries)
print('\n'.join(map(str, ans)))
