AtCoder Regular Contest 117 B,C,D問題メモ

AtCoder Regular Contest 117

ちょっとずつ変なケアレスミスでバグらせてる時間が勿体ないねえ。
春は眠い季節だからミスも多くなるの。¹⁾

B - ARC Wrecker

• $N$ 棟のビルが並んでいて、$i$ 番目の高さは $A_i$
• 「ある数字 $X$ を指定して、その時点で高さが $X$ 以上の全てのビルの高さを1ずつ低くする」操作を繰り返せる
  • ダルマ落としでダルマがいっぱいあるイメージ
• 操作を0回以上行った後の景観（各ビルの高さの並び）としてあり得るものの個数を $\mod{10^9+7}$ で求めよ
• $1 \le N \le 10^5$
• $1 \le A_i \le 10^9$

各ビルの最終的な高さを $B$ とすると、最終的な景観としてあり得ないのは、

• 最初、$A_i \lt A_j$ だった $(i,j)$ が、最終的には $B_i \gt B_j$ となっている
• 最初、$A_i - A_j = 5$ だった $(i,j)$ が、最終的には $B_i - B_j = 6$ などと差が増えている

このような $i,j$ があったらダメ。

ビルは並び順は関係ないので、昇順にソートしておき、$i$ も振りなおす。
以下のDPを考える。

• $DP[i][j]=i$ 番目のビルまでみて、最終的に残った一番高いビル（つまりビル $i$）の高さが $j$ の時のパターン数

すると、たとえば $A=(1,2,5,5,7)$ において、

Ai=1

j
7                           1
6                           4
5                 1    1    8
4                 3    3   11
3              ┌ 4    4   11
2            1 ├ 4    4 ┌ 8
1       1 ┌ 2 ├ 3    3 ├ 4
0   1   1 ┴ 1 ┴ 1 ─ 1 ┴ 1
------------------------------
 i (0)  1    2    3    4    5
Ai      1    2    5    5    7

このようになる。

$A_{i+1}$ と $A_i$ の差がたとえば $3$ だったとして、$DP[i][j]$ は、$DP[i+1][j]～DP[i+1][j+3]$ の4箇所にそれぞれ1倍で遷移する。 この範囲に遷移する限り、最終的な景観があり得る状態に保たれる。

さて、このDPを愚直に計算すると $O(N A_{max})$ だが、どの遷移も「$A_{i+1}-A_i+1$ 箇所に1倍で遷移する」というのが変わらない。

つまり、最終的に必要なのは総和なので、最初から総和で管理して問題ない。

答えは、$(A_1+1) \times (A_2-A_1+1) \times (A_3-A_2+1) \times ...$ となる。

Python3

import sys

n, *aaa = map(int, sys.stdin.buffer.read().split())
MOD = 10 ** 9 + 7
aaa.sort()
p = 0
ans = 1
for a in aaa:
    ans = ans * (a - p + 1) % MOD
    p = a

print(ans)

C - Tricolor Pyramid

• 青、白、赤の3色の積み木が無限にあり、これでピラミッドを作る
• 今、最下段の色が決まっている
• 以下のルールで積み上げていく
  • 色 $c_1,c_2$ と並んだ部分の上にのせる積み木の色は、
    • $c_1=c_2$ の場合、同じ色
    • $c_1 \neq c_2$ の場合、そのどちらでもない1色
• 一番上の積み木の色を求めよ
• $2 \le 最下段の長さ \le 4 \times 10^5$

取り得る値が3つの列をいろいろ操作する問題は、「操作しても変わらない不変量」を見つけるのが解法であることが多い。

今回の場合、足したり引いたりXORしたりして試すと、

• B,W,R をそれぞれ 0,1,2 で置き換える
• $c_1,c_2$ と並んだ部分の上の色は、$-c_1-c_2 \mod{3}$

という法則が見つかる。それをもって下から計算すると、

                  c1+4c2+6c3+4c4+c5
           -c1-3c2-3c3-c4   -c2-3c2-3c4-c5
        c1+2c2+c3     c2+2c3+c4     c3+2c4+c5
  -c1-c2        -c2-c3        -c3-c4        -c4-c5
c1         c2            c3            c4            c5

このように、一番下の段の値に二項係数をそれぞれかけたものが一番上の段の値となる。

この問題はそこからもう一段階ある。

二項係数は $N=4 \times 10^5$ にもなると値が非常に大きくなってやってらんないので、適当にmodをとる必要があるが、 そのときの法は「計算過程で出てくる値のどれとも互いに素な数」で無いと適切に計算できない。

例えば $\mod{1000000007}$ が $4$ になったとしても、可能性のある値は $4,1000000011,2000000018,...$ となり、 それぞれの $\mod{3}$ は $1,0,2,...$ なのでどれなのか結局わかんない。

調べると、任意modでの二項係数という記事が出てくる。

• 任意 mod で二項係数を列挙する - suisen のブログ

これによると、素因数分解した形で管理するといいらしい。

今回は $3$ という非常に単純な値なので、「3で割ったときに 0,1,2 となる素因数をそれぞれいくつ持つか」で管理してみる。

     0  1  2
 1   0  0  0
 2   0  0  2
 6   1  0  1  (2*3)
 7   0  1  0  (7)
10   0  0  2  (2*5)

（公式解説によると、実際は「3で何回割れるか」と「その後に残る値の $\mod{3}$」の2つで管理できるらしい）

そして、二項係数 ${}_NC_r$ を順番に求める場合、以下のようにすると順次計算できる。

• ${}_NC_0=1$
• ${}_NC_i={}_NC_{i-1} \times \dfrac{N-i+1}{i}$

(例)
     i    0      1     2     3  ...
  ---------------------------------
                10     9     8
10_C_i    1  x  --  x  -  x  -
                 1     2     3

今回は、これに従って素因数分解結果(mod3)を、かけ算なら足し合わせ、割り算なら引けば、二項係数 $\mod{3}$ を表現できる。

3つの値をそれぞれ $d_0,d_1,d_2$ としたとき、その値の $\mod{3}$ は、

• $d_0 \gt 0$（つまり $3$ を素因数に1つでも持つ）場合、$0$
• それ以外の場合、$2^{d2}$

これに従って、最下段の並びに、二項係数mod3をかけて足し合わせたものが、最上段の色を示す値となる。

ただし、最下段の長さが偶数の場合、正負逆転した値となることに注意。

より汎用的な 素数modの二項係数の求め方としては、Lucasの定理というのがある。

• Lucasの定理とその証明 | 高校数学の美しい物語

• 二項係数$\mod{M}$ は、$M$ によって、以下の方法でできる、ということか
  • 求めたい ${}_nC_r$ のどの $n,r$ より大きい素数 … 階乗$\mod{M}$ を事前計算（いつもの）
  • それ以外の素数 … Lucasの定理
  • 素数に限らない … 任意 mod で二項係数を列挙する - suisen のブログ の方法

Python3

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=None):
    """ xを素因数分解、辞書型で返す """
    # primes: sqrt(x)までの素数リスト
    #   省略時は奇数を試し割り（2だけ別処理）
    #   単発なら不要だが、複数のxで行う場合は事前計算した方がよい

    factors = [0, 0, 0]
    for p in primes:
        if p * p > x:
            break
        while x % p == 0:
            factors[p % 3] += 1
            x //= p
    if x != 1:
        factors[x % 3] += 1
    return factors


# 愚直解法の実験
# import random
#
# ccc = random.choices(range(1, 4), k=20)
# print(ccc)
# for _ in range(len(ccc) - 1):
#     ddd = []
#     for c1, c2 in zip(ccc, ccc[1:]):
#         if c1 == c2:
#             ddd.append(c1)
#         else:
#             ddd.append(c1 ^ c2)
#     ccc = ddd
#     print(ccc)


n = int(input())
ccc = input()
ccc = list(map('BWR'.index, ccc))
primes = get_primes(400000)
pfs = [[0, 0, 0], [0, 0, 0]] + [prime_factorize(x, primes) for x in range(2, n + 1)]
check = [0, 0, 0]
case = ccc[0]

for i, c in enumerate(ccc[1:], start=1):
    pf1 = pfs[n - i]
    pf2 = pfs[i]
    check[0] += pf1[0] - pf2[0]
    check[1] += pf1[1] - pf2[1]
    check[2] += pf1[2] - pf2[2]

    if check[0] > 0:
        continue
    case = (case + c * pow(2, check[2], 3)) % 3

if n % 2 == 0:
    case = -case % 3
   
print('BWR'[case])

mod3の不変量を求める方針は一緒。

二項係数 ${}_nC_r$ のmod3を計算していくと、$n=3,9,27,...$ といった3のべき乗で、 '1 0 0 0 … 0 1' のように ${}_nC_0={}_nC_n=1$ を除いて全て間が0となることを利用する。

つまり、

• $c1,c2,c3,c4$ という並びで操作を3回繰り返すと、最後に残るのは $-(c1+c4)$ となる
• $c1,c2,...,c{10}$ という並びで操作を9回繰り返すと、最後に残るのは $-(c1+c{10})$ となる
• $c1,c2,...,c{28}$ という並びで操作を27回繰り返すと、最後に残るのは $-(c1+c{28})$ となる

なので、$3^k$ 段上の状態を、一足飛びに計算してしまえる。

この操作は、最下段の長さを $N$ として $O(\log_3{N})$ 回程度ですみ、1回の操作もどんどん短くなるので、全体で $O(N)$ 程度で計算できる。

実装も短い。

Python3

n = int(input())
ccc = list(map('BWR'.index, input()))
l = len(ccc) - 1

while l > 2:
    k = 3
    while k * 3 <= l:
        k *= 3
    ccc = [(-c0 - ck) % 3 for c0, ck in zip(ccc, ccc[k:])]
    l -= k

if l == 0:
    ans = ccc[0]
elif l == 1:
    ans = -sum(ccc) % 3
else:
    ans = (sum(ccc) + ccc[1]) % 3

print('BWR'[ans])

D - Miracle Tree

• $N$ 頂点の木が与えられる
• 各頂点へ、整数を1つずつ書き込む（頂点 $i$ に書き込む値を $E_i$ とする）
• 以下の条件を満たす書き込み方を1つ構築せよ
  • 全て正整数（$E_i \ge 1$）
  • 全ての $i,j$ について、$|E_i-E_j| \ge 木における(i,j)の距離$
  • 上記の2つを満たす中で、$E_i$ の最大値が最小
• $2 \le N \le 2 \times 10^5$

一直線の木なら簡単で、1から順に埋めていくだけ。

①--②--③--④--⑤

二股に分かれていても、基本、一直線に取れる部分はこの法則より小さくできない。

①--②--③--④--⑤
        └--○--○

で、残る部分が別の一直線に含まれるように変形させると、

            ↓~~~~~~~~⑤とのペアを考えると、5との差が3以上必要
⑤--④--③--○--○  ←⑤とのペアを考えると、5との差が4以上必要
        └--②--①

そのため、以下のようになる。5との差が3以上必要といっても、減らしてしまうと今度は②などとの整合が取れなくなるので、増やす方となる。

⑤--④--③--⑧--⑨
        └--②--①

これは、DFSの訪問順で表現できる。

つまり、根を決めて、$k=1$ から、隣接する頂点に移動するたびに（訪問済みであろうと）$k$ を1ずつ増やしてDFSすると、各頂点の値は「その頂点に初めて訪れられたときの $k$ の値」とすればよい。

①--②--③--④--❺  k=5まで
        └--○--○

①--②--③--❹--⑤  k=6
        └--○--○

①--②--❸--④--⑤  k=7
        └--○--○

①--②--③--④--⑤  k=8
        └--❽--○

①--②--③--④--⑤  k=9
        └--⑧--❾

さて、ではその中で「最大値を最小化」するためにはどんな条件が必要か。

どの頂点を根に選ぼうが「根からDFSで全頂点を訪れ根に戻ってくる時の $k$ の値」（オイラーツアーの長さ）は変わらない。

根が中途半端な位置だった場合、根を①とするより、一番最初にたどり着いた葉を①とした方がいいのは明らかで、「根から、最初にたどり着いた葉までの $k$ の増分」は省略できる。

また、実際に頂点の値として使われるのは最後にたどり着いた葉なので、「最後にたどり着いた葉から、根に戻るまでの $k$ の増分」も省略できる。

これが最大になるように、最初と最後にたどり着く葉を選べばよい。

つまり、木の直径となる。

• BFSを2回行い、木の直径の両端となるような2頂点 $u,v$ を調べる
• $u$ を根として、$v$ を最後に訪れるようにDFSする
• すると、上記の番号の振り方で、最大値を最小化できる

Python3

import sys

from collections import deque

sys.setrecursionlimit(400005)


def bfs(s, links):
    q = deque()
    q.append(s)
    stacked = {s}
    parents = [-1] * len(links)
    v = -1
    while q:
        v = q.popleft()
        for u in links[v]:
            if u in stacked:
                continue
            q.append(u)
            stacked.add(u)
            parents[u] = v
    return v, parents


n = int(input())
links = [set() for _ in range(n)]
for _ in range(n - 1):
    a, b = map(int, input().split())
    a -= 1
    b -= 1
    links[a].add(b)
    links[b].add(a)

v1, _ = bfs(0, links)
v2, parents = bfs(v1, links)

# v1を始点としたときにv2までのパス上にある頂点を調べておく
ddd = [False] * n
p = v2
while p != v1:
    ddd[p] = True
    p = parents[p]

ans = [0] * n


def dfs(k, v):
    ans[v] = k
    keep = -1

    for u in links[v]:
        if u == parents[v]:
            continue
        if ddd[u]:
            keep = u
            continue
        k = dfs(k + 1, u) + 1

    if keep != -1:
        k = dfs(k + 1, keep) + 1

    return k


dfs(1, v1)

print(*ans)