AtCoder Beginner Contest 482 E,F問題メモ
E - Diameter Endpoint Colors
問題文
$N$ 頂点の木 $T$ が与えられます。
$i=1,2,\dots,N-1$ について、$i$ 番目の辺は頂点 $u_i$ と頂点 $v_i$ を結んでいます。
各頂点には色が塗られており、頂点 $v$ は色 $C_v$ で塗られています。
$1 \leq u \lt v \leq N$ を満たす整数組 $(u,v)$ で以下の条件を全て満たすものの個数を求めてください。
ただしパスの長さはそのパスに含まれる辺の本数であり、木の直径はその木に含まれる任意の単純パスの長さの最大値です。
制約
解法
「木の直径・中心」を知っていますかという問題。これまでにもABCでいくつか類題は出題されている。
木の中心が1個の場合
⓪
,--'｜`--,
① ② ③
： ： ：
： ⑥⑦ ：
④⑤ ⑧⑨
中心⓪を根とした時の⓪に隣接する頂点（①②③）の部分木それぞれで、
最遠頂点までの距離と、最遠となる頂点集合をそれぞれ求める。
木の直径（中心が1個の場合は必ず偶数）の半分を $h$ として、⓪と最遠頂点との距離が $h$ 未満の部分木は除外。
$h$ である部分木だけ残し、異なる2つの部分木の最遠頂点から1つずつ選んだ同じ色の頂点ペアの個数が答え。
木の中心が2個の場合
こっちの方が楽。
①-----②
： ：
： ：
④⑤ ⑥⑦
2つの中心をそれぞれ根として部分木を1個の場合と同様に求める。
2つの部分木の最遠頂点から1つずつ選んだ同じ色の頂点ペアの個数が答え。
Python3
from collections import deque, Counter
def get_tree_center(links):
"""
木の直径と、中心の頂点番号を求める（1辺の距離を1とみなす）
:param links: links[v] = {u1, u2, ...} 隣接する頂点番号
:return: (Diameter, c1, c2) 直径が偶数の場合、中心は単一頂点のためc1が中心、c2=-1。奇数の場合、中心はc1-c2を結ぶ辺
"""
n = len(links)
# first bfs
q = deque([0])
dists = [-2] * n
dists[0] = -1
v = -1
while q:
v = q.popleft()
for u in links[v]:
if dists[u] != -2:
continue
q.append(u)
dists[u] = -1
v1 = v
# second bfs
q = deque()
q.append(v1)
predecessors = [-1] * n
dists[v1] = 0
v = -1
while q:
v = q.popleft()
for u in links[v]:
if dists[u] != -1:
continue
q.append(u)
dists[u] = dists[v] + 1
predecessors[u] = v
v2 = v
d = dists[v]
if d % 2 == 0:
hd = d // 2
i = v2
while i != v1:
if dists[i] == hd:
return d, i, -1
i = predecessors[i]
else:
hd = (d + 1) // 2
i = v2
while i != v1:
if dists[i] == hd:
return d, predecessors[i], i
i = predecessors[i]
raise AssertionError('Not Reachable')
def solve(n, edges, ccc):
links = [[] for _ in range(n)]
for u, v in edges:
u -= 1
v -= 1
links[u].append(v)
links[v].append(u)
diameter, c1, c2 = get_tree_center(links)
INF = 1 << 60
def get_distance(c1, c2, dists):
q = deque([c1])
dists[c1] = 0
max_d = 0
same_dists = [c1]
while q:
u = q.popleft()
for v in links[u]:
if u == c1 and v == c2:
continue
if dists[v] == INF:
dists[v] = dists[u] + 1
if dists[v] > max_d:
max_d = dists[v]
same_dists = [v]
elif dists[v] == max_d:
same_dists.append(v)
q.append(v)
return max_d, same_dists
ans = 0
dists = [INF] * n
if c2 == -1:
h = diameter // 2 - 1
available_subtrees = []
for ch in links[c1]:
md, sd = get_distance(ch, c1, dists)
if h == md:
available_subtrees.append(sd)
cnt = Counter()
for uuu in available_subtrees:
new = Counter()
for u in uuu:
c = ccc[u]
ans += cnt[c]
new[c] += 1
for c, a in new.items():
cnt[c] += a
else:
_, uuu = get_distance(c1, c2, dists)
_, vvv = get_distance(c2, c1, dists)
cnt1 = Counter(ccc[u] for u in uuu)
cnt2 = Counter(ccc[v] for v in vvv)
for c, a in cnt1.items():
if c in cnt2:
ans += a * cnt2[c]
return ans
n = int(input())
edges = [list(map(int, input().split())) for _ in range(n - 1)]
ccc = list(map(int, input().split()))
ans = solve(n, edges, ccc)
print(ans)
F - Mod M Divisor
問題文
制約
解法
$M=1$ の場合は明らかに $k=0$ に対して $t=1$ が答え。以下、$M \ge 2$ とする。
TLE 解法
$N$ の約数集合を $D=(d_1,d_2,...,d_k)$ とする。
以下の $M$ 頂点のグラフを考える。
頂点に $0,1,...,M-1$ の番号が付いている。
頂点 $v$ からは、$vd_1 \pmod{M}, vd_2 \pmod{M},..., vd_k \pmod{M}$ に有向辺がある。
この時、以下の $A=(A_0,...,A_{M-1})$ が答えとなる。
だが、$10^{14}$ 以下の約数は最大 $\sigma=17280$ 個あるので、$O(\sigma M) \simeq 10^9$ となり間に合わない。
高速化
辺を約数集合 $D$ と対応づけると、
グラフの最短距離に直接結びつけてイメージしやすいのだが、辺の個数が増えてしまう。
その前段、素因数分解の結果に対応づけるだけで、高速化に繋がる。
$N$ の素因数分解を $P=(p_1,p_2,...,p_k)$ とする。
$N=12$ なら $P=(2,2,3)$ のように、同じ素因数も展開する。制約の範囲では $k \le 46$ 程度が保証される。
グラフの遷移を、大きな「フェーズ」と小さな「ステップ」に分ける。
フェーズが、TLE解法の1回の遷移「キュー内の各頂点から約数集合 $D$ の各辺で1回遷移する処理」に近く、
ステップが「その遷移を素因数ごとに分解して行う処理」となる。
こうすることで、$M$ 個の各頂点からの遷移が最大 $2k$ 程度に抑えられる。
（フェーズの最初の方に確定した頂点は、現フェーズ+次フェーズでそれぞれ約 $k$ の遷移が走る。最大 $k$ にすることも可）
AIに図化させた遷移図。用語が上記説明と違っていたり、ところどころ実線と点線が変な気がするが、まぁ言葉よりは断然わかりやすい。
Python3
from collections import defaultdict, deque
def get_primes(n):
""" nまでの素数リスト """
candidates = [0, 1] * (n // 2 + 1)
if n % 2 == 0:
candidates.pop()
candidates[1] = 0
candidates[2] = 1
for p in range(3, int(n ** 0.5) + 1, 2):
if candidates[p]:
for m in range(p * p, n + 1, 2 * p):
candidates[m] = 0
return [p for p, b in enumerate(candidates) if b]
def prime_factorize(x, primes):
""" xを素因数分解、辞書型で返す """
# primes: sqrt(x)までの素数リスト
# 省略時は奇数を試し割り（2だけ別処理）
# 単発なら不要だが、複数のxで行う場合は事前計算した方がよい
def bit_length(n):
x = 0
while 1 << x < n:
x += 1
return x
factors = defaultdict(int)
for p in primes:
if p * p > x:
break
if x % p == 0:
factors[p] = 1
x //= p
while x % p == 0:
factors[p] += 1
x //= p
if x != 1:
factors[x] = 1
return factors
def get_divisors(pf):
ppp = sorted(pf.keys())
divisors = [1]
for p in ppp:
cnt = pf[p]
w = len(divisors)
x = p
for c in range(cnt):
for i in range(w):
a = divisors[i]
divisors.append(a * x)
x *= p
return divisors
def solve(n, m):
if m == 1:
return [1]
primes = get_primes(10 ** 7 + 1)
pf = prime_factorize(n, primes)
links = []
for p, c in pf.items():
links.extend([p % m] * c)
ans = [-1] * m
ans[1] = 1
q = [1]
for i in range(1, m + 1):
nxt = []
for l in links:
lim = len(q)
for j in range(lim):
v = q[j] * l % m
if ans[v] == -1:
ans[v] = i
q.append(v)
nxt.append(v)
q = nxt
if len(q) == 0:
break
return ans
n, m = list(map(int, input().split()))
ans = solve(n, m)
print('\n'.join(map(str, ans)))
