AtCoder Regular Contest 116 A,B,C,D,E,F問題メモ

AtCoder Regular Contest 116

Dのイージーミスの原因究明までにWA出しまくったのがもったいなかった。。。

A - Odd vs Even

• 正整数 $N$ の約数には、偶数が多いか奇数が多いか同じか、判定せよ
• テストケースは $T$ 個与えられる
• $1 \le T \le 2 \times 10^5$
• $1 \le N \le 10^{18}$

素因数分解したときの2の個数で判定できる。

素因数のうち、2以外（奇数だけ）を組み合わせて作られる約数は全て奇数。

そこに2が1個あったら、それぞれにかけることで、同数の偶数の約数が作られる。

N =       3^2 * 5  →  1  3  5  9 15  45
N = 2   * 3^2 * 5  →  2  6 10 18 30  90 （上記に加えて）
N = 2^2 * 3^2 * 5  →  4 12 20 36 60 180 （上記にさらに加えて）

2が2個あったら、奇数の2倍の偶数の約数が作られる。この時点で奇数は逆転不可能。

よって、$N$ が奇数なら'Odd'、2の倍数で4の倍数でないなら'Same'、4の倍数なら'Even'となる。

Python3

t = int(input())
for _ in range(t):
    n = int(input())
    if n % 2 == 1:
        print('Odd')
    elif n % 4 == 0:
        print('Even')
    else:
        print('Same')

B - Products of Min-Max

• 長さ $N$ の整数列 $A$ が与えられる
• $A$ の空でない部分列 $B$ は $2^{N-1}$ 個あるが、全てについて $\max(B) \times min(B)$ を計算し、その総和を $\mod{998244353}$ で求めよ
• $1 \le N \le 2 \times 10^5$

この問題における部分列は、連続していなくてもいい点に注意。そのため、最初に $A$ をソートして考えて問題ない。

最小値,最大値を固定し、これらが $\min(B),\max(B)$ になるような $B$ が何個あるか考える。

 max=   2  3  5  6  8  9

min=2   1  1  2  4  8 16
min=3      1  1  2  4  8
min=5         1  1  2  4
min=6            1  1  2
min=8               1  1
min=9                  1

例えば $\min=2,\max=6$ の時、間の $3,5$ が入るか入らないかで $2^2=4$ 通りの $B$ があり、答えに寄与するのは $2 \times 6 \times 4=48$ となる。

最小値を中心にまとめると、$\min=2$ のとき、$2 \cdot 1 + 3 \cdot 1 + 5 \cdot 2 + ... + 9 \cdot 16$ だけ寄与する。

表を見ると、minが1つ大きい値になるたびに、各値とのペアの数は概ね $\dfrac{1}{2}$ ずつになっていく。

なので、最初に $\min=2$ の時の係数 $C$ を上記の通り求めると、 次の $\min=3$ の時の係数は、$C'=\dfrac{C - 2 + 3}{2}$ で求められる。

さらに次の $\min=5$ の時の係数は $C''=\dfrac{C'-3+5}{2}$ ……と、逐次的に $O(1)$ で求められる。

Python3

n = int(input())
aaa = list(map(int, input().split()))
MOD = 998244353
INV2 = pow(2, MOD - 2, MOD)
aaa.sort()
coef = aaa[0]
pow2 = 1
for a in aaa[1:]:
    coef = (coef + a * pow2) % MOD
    pow2 = pow2 * 2 % MOD

ans = 0
for i, a in enumerate(aaa):
    ans = (ans + a * coef) % MOD
    if i < n - 1:
        coef = (coef - a + aaa[i + 1]) * INV2 % MOD

print(ans)

C - Multiple Sequences

• 正整数 $N,M$ が与えられる
• 長さ $N$ の整数列 $A$ で、以下の条件をともに満たすものの数を $\mod{998244353}$ で求めよ
• 条件
  • $1 \le A_i \le M$
  • $A_{i+1}$ は $A_i$ の倍数
• $1 \le N,M \le 2 \times 10^5$

$A_N$ の値を $1～M$ まで全探索する。

$A_N=k$ と固定した時、$k$ の素因数分解がたとえば $360=2^3 \ 3^2 \ 5$ とかだったとすると、

• $A_i$ と比べて $A_{i+1}$ が2倍になる箇所が、重複を許して3箇所
• $A_i$ と比べて $A_{i+1}$ が3倍になる箇所が、重複を許して2箇所
• $A_i$ と比べて $A_{i+1}$ が5倍になる箇所が、重複を許して1箇所

となるので、そのような数列の個数は重複組み合わせを使って ${}_NH_3 \times {}_NH_2 \times {}_NH_1$ と数えられる。

便宜的に $A_0=1$ とすると、○倍とする箇所の候補は $N$ 箇所というのがわかる。

  A0  A1  A2  A3  ...   AN
   1   ?   ?   ?       360
    ↑  ↑  ↑  ↑   ↑    ←N箇所

計算量は、素数を事前計算しておけば

• $M$ 個のそれぞれにつき
  • $O(\log{M})$ で素因数分解
  • $O(素因数の個数)$ 回、重複組み合わせを計算して掛け合わせる

で計算できる。素因数の個数は最大でもそんなに大きくならないし、ほとんどは1～3個程度なのでまぁ定数と見てよく（乱暴）、$O(M \log{M})$ 程度となる。

Python3

import os
import sys

import numpy as np


def solve(n, m):
    def mod_pow(x, a, MOD):
        ret = 1
        cur = x
        while a:
            if a & 1:
                ret = ret * cur % MOD
            cur = cur * cur % MOD
            a >>= 1
        return ret

    def prepare_factorials(n, MOD):
        factorials = np.ones(n + 1, np.int64)
        for m in range(1, n + 1):
            factorials[m] = factorials[m - 1] * m % MOD
        inversions = np.ones(n + 1, np.int64)
        inversions[n] = mod_pow(factorials[n], MOD - 2, MOD)
        for m in range(n, 1, -1):
            inversions[m - 1] = inversions[m] * m % MOD
        return factorials, inversions

    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):
        factors = {}
        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

    MOD = 998244353

    facts, finvs = prepare_factorials(n + m, MOD)
    primes = get_primes(int(m ** 0.5) + 5)

    ans = 1
    for k in range(2, m + 1):
        factors = prime_factorize(k, primes)
        tmp = 1
        for c in factors.values():
            tmp = tmp * facts[n + c - 1] % MOD
            tmp = tmp * finvs[c] % MOD
            tmp = tmp * finvs[n - 1] % MOD
        ans = (ans + tmp) % MOD

    return ans


SIGNATURE = '(i8,i8)'
if sys.argv[-1] == 'ONLINE_JUDGE':
    from numba.pycc import CC

    cc = CC('my_module')
    cc.export('solve', SIGNATURE)(solve)
    cc.compile()
    exit()

if os.name == 'posix':
    # noinspection PyUnresolvedReferences
    from my_module import solve
else:
    from numba import njit

    solve = njit(SIGNATURE, cache=True)(solve)
    print('compiled', file=sys.stderr)

n, m = np.fromstring(sys.stdin.read(), dtype=np.int64, sep=' ')
ans = solve(n, m)
print(ans)

D - I Wanna Win The Game

• 正整数 $N,M$ が与えられる
• 長さ $N$ の整数列 $A$ で、以下の条件を全て満たすものの数を $\mod{998244353}$ で答えよ
• 条件
  • 全要素は非負
  • 全要素の和は $M$
  • 全要素のXORは $0$
• $1 \le N,M \le 5000$

なんかC問題と似てる。解き方もまぁ似てるような気がする。

まず例外処理として、$N=1$ の時、XORを0にできず不可能なので除く。

総XOR=0の制約から、各bitについて、そこを'1'とする項は偶数個とならなければいけない。

なので、上位bitから、そのbitを'1'にする項の個数を「偶数」「個数が $N$ を超えない」「より上位のbitの結果と併せて、総和が $M$ を超えない」範囲で遷移するDPを作る。

• $DP[i][j]=$ 上位から $i$ 番目のbitまでで立てるbitを決め終わって、総和を $M$ とするための残りが $j$ である数列の個数

初期値は $DP[0][M]=1$ となる。

N = 5   M = 20 = 10100(2)

bit    DPの内容
       10100:  1    初期状態

 1000  10100:  1    
         100: 10    残り10100からbit'1000'を2個立てると残り100,
                    立て方の場合の数は二項係数で 1 x 5C2通り
  
  100  10100:  1
         100: 15    残り10100からbitを4個, 個数  1 x 5C4 を加算
        1100: 10    残り10100からbitを2個, 個数  1 x 5C2

   10  10100:  1
         100: 65    残り 1100からbitを4個, 個数 10 x 5C4 加算
        1100: 15    残り10100からbitを4個, 個数  1 x 5C4 加算
       10000: 10    残り10100からbitを2個, 個数  1 x 5C2
        1000:100    残り 1100からbitを2個, 個数 10 x 5C2
           0:150    残り  100からbitを2個, 個数 15 x 5C2

このように、残り $j$ からbit $b$ を $k$ 個立てた時、$DP[i][j-bk]+=DP[i-1][j] \times {}_NC_k$ で遷移できる。

最終的に、$DP[Mのbit長][0]$ が答え。

計算量は、DPにおける $i$ の状態数が $O(\log{M})$、$j$ の状態数が $O(M)$、さらに各 $DP[i][j]$ につき何個のbitを立てるかで $O(N)$。あわせて $O(NM\log{M})$ となる。

もう少し細かく見ると、$j$ の状態数はdictなどで持つと $i$ が進む毎に1,2,4,8,…と2倍ずつになっていくので、$(i,j)$ の組の状態数がまとめて $O(M)$ となり、全体で $O(NM)$ となる。

Python3

from collections import defaultdict


def prepare(n, MOD):
    f = 1
    factorials = [1]
    for m in range(1, n + 1):
        f = f * m % MOD
        factorials.append(f)
    f = pow(f, MOD - 2, MOD)
    finvs = [1] * (n + 1)
    finvs[n] = f
    for m in range(n, 1, -1):
        f = f * m % MOD
        finvs[m - 1] = f
    return factorials, finvs


def solve(n, m):
    if n == 1 or m % 2 == 1:
        return 0

    MOD = 998244353
    facts, finvs = prepare(n, MOD)

    dp = defaultdict(int)
    dp[m] = 1
    d = 1 << m.bit_length()

    facts2 = [0] * (n + 5)
    for c in range(0, n + 1, 2):
        facts2[c] = facts[n] * finvs[c] * finvs[n - c] % MOD

    while d:
        prev_dp = list(dp.items())
        for c in range(2, n + 1, 2):
            dc = d * c
            if dc > m:
                break
            for key, value in prev_dp:
                if dc > key:
                    continue
                dp[key - dc] = (dp[key - dc] + facts2[c] * value) % MOD
        d >>= 1

    return dp[0]


n, m = map(int, input().split())

print(solve(n, m))

E - Spread of Information

• $N$ 頂点の木と見なせる都市と道路がある
• 国王が、時刻 $0$ に $K$ 個の都市に情報を伝える
• 時刻 $t$ に情報が伝えられた各都市は、時刻 $t+1$ に、自身と道路で直接つながっている全ての都市に情報を伝える
• 適切に最初の $K$ 個の都市を選んだとき、全ての都市に情報が伝わるのにかかる最小時間を求めよ
• $1 \le K \lt N \le 2 \times 10^5$

答えを決め打つ二分探索。

国王から時刻0に情報を伝えられる都市を、初期都市と呼ぶことにする。

時間 $m$ を決め打って、$m$ 以内に全都市に行き渡らせるのが可能かどうかは、木DPをおこなうことで判定できる。

適当に根を1つ決め、根付き木としておく。

• $DP[v]=$ 頂点 $v$ 以下の部分木について、適切に初期都市を決めた時の、以下のいずれか
  • 保留中の頂点のうち最も深いものが、頂点 $v$ から見て深さ○○の位置にある（$v$ の深さを1とする）
    • 言い換えると、「このカウントが $m$ を超える前に、祖先に初期都市を置く必要がある」
  • 既に $v$ 以下に置いた初期都市によって、祖先のうち距離○○まではカバーできる

2つは表裏一体なので、例えば前者を正、後者を負でもっておけばよい。

保留中の頂点とは、「まだ部分木内には自身をカバーできる初期都市が無いが、祖先に初期都市を置けばカバーが間に合う頂点」とする。 基本的に、保留中の頂点をカバーできなくなる前ギリギリで初期都市を置いていくことになる。

わかりやすいように1本道グラフを考えると、葉は $DP[v]=1$ となり、親をたどる毎に $2,3,4,...$ と増える。 $m$ になったらその次は初期都市とする。初期都市は $DP[v]=-m$ となる。

m=2
   根
   ○--●--○--○--○--○--●--○--○    ●:初期都市
DP -1  -2   2   1   0  -1  -2   2   1

複数の子を持つ場合、その統合は少し考える必要がある。

$v$ の子頂点の中で $DP$ の正の最大値 $a$ と負の最小値 $b$ を得る。

a+1 < -b のとき

    ⓥ:-2
  ／｜＼     bの初期都市によって、他の子孫の保留中の頂点はすべてカバーできる
 -3  2  1
 |   |       DP[v] = b+1 となる
...  1

a = m のとき

    ⓥ:-3  m=3
  ／｜＼     vを初期都市とする必要がある
 -2  3  1    
 |   |       DP[v] = -m となる
... ...

それ以外

    ⓥ:3     bの初期都市ではaの保留中の頂点に届かないため、
  ／｜＼     aの方の制約を優先する必要がある
 -2  2  1    
 |   |       DP[v] = a+1 となる
...  1

全頂点のDPを埋めたとき、$DP[根]$ が正なら、初期都市をもう1つ置く必要がある。（たとえば根頂点に置けば問題なくカバーできる）

これが $m$ を決めたときの最小ギリギリの初期都市の置き方となり、$K$ を超えないかどうかで判定できる。

1回の判定に $O(N)$、全体で $O(N \log{N})$ で求められる。

Python3

def check(dfs_order, links, n, m, k):
    remaining = [0] * n
    count = 0
    for v in dfs_order:
        c = 0
        r = 0
        for u in links[v]:
            ru = remaining[u]
            if ru >= 0:
                c = max(c, ru)
            else:
                r = min(r, ru)
        if c + 1 <= -r:
            remaining[v] = r + 1
        elif c == m:
            remaining[v] = -m
            count += 1
            if count > k:
                return False
        else:
            remaining[v] = c + 1

    if remaining[0] > 0:
        count += 1
        if count > k:
            return False

    return True


def solve(n, k, links):
    dfs_order = []
    parents = [-1] * n
    status = [0] * n
    q = [0]
    while q:
        v = q[-1]
        if status[v] == 0:
            for u in links[v]:
                if u == parents[v]:
                    continue
                q.append(u)
                parents[u] = v
                links[u].remove(v)
            status[v] = 1
        else:
            q.pop()
            dfs_order.append(v)

    l = 0
    r = n
    while l + 1 < r:
        m = (l + r) // 2
        if check(dfs_order, links, n, m, k):
            r = m
        else:
            l = m
    return r


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

print(solve(n, k, links))

F - Deque Game

• $K$ 個の数列が与えられる
• AとBの2人でゲームを行う
• ルール
  • 交互に、長さが2以上の数列を1つ選び、その先頭または末尾要素を削除する
  • 全ての数列の長さが1になったら終了
  • Aが先手
• Aは最後に残る $K$ 個の要素の和を最大化、Bは最小化したいと考えている
• 両者最適に行動するとき、最後に残る $K$ 個の要素の和を求めよ
• $1 \le 全数列の長さ合計 \le 2 \times 10^5$
• $1 \le A_{i,j} \le 10^9$

ゲーム問題。

数列が1個しかなく、長さが短い場合から考える。
ただしAが先手とは限らず、どちらから始める場合も考える。

操作の余地がない。その1項が答え。

先に操作した方が好きなのを残せる。

この辺からちょっとややこしくなって、

Aから手を着ける場合
         A手番後  B手番後
1 2 3  →  2 3  →  2
1 3 2  →  3 2  →  2
2 1 3  →  1 3  →  1

Bから手を着ける場合
         B手番後  A手番後
1 2 3  →  1 2  →  2
1 3 2  →  1 3  →  3
2 1 3  →  2 1  →  2

このように、先に手を着けた側が一番残したくないもの （Aにとっては最小値、Bにとっては最大値）が中央にある場合のみ、 それが残ってしまう。

その他の場合は、中央値が残る。

左端と右端のどちらを取り除いた方がよいかは、3個のときに帰着して両方調べれば確定する。
先手が得する方を選ぶ。

先手だった方が、残り3個になったときにも先手となる。

真ん中3つの要素が重要となる。結論から言うと、真ん中の3つから始めたときと結果は同じになる。

      v v v
... 9 3 5 7 9 ...

Aが先手の時、ちょっと欲張れば9とかに手が届きそうに見えるが、 BはAと逆方向をとり続けることで真ん中を3,5,7の状態に保ち続けることができ、9にはできない。

またBにとっては、ひょっとしたら3,5,7を中央に保つより、 中央をずらすことで残す値をもっと小さくできるかも知れない。 しかし、その場合はAが阻止できる。

          v v v
1 1 ... 1 3 5 7 1 ... 1 1

Bとしては1を残すべく、中央をずらしたいところである。
Aがたとえば左を取り、Bも同じ側を取ると、中央が5,7,1にずれる。
しかし、そしたらAは次からは右しか取らないことで、これ以上、中央がずれるのを阻止できる。

このようにずらして中央が 5 7 ? となった際にこの3数から残される値は、?の値によらず5以上となる。

もしAが右を取ったら、ずらされた中央 ? 3 5 から残される値は3以上なので、 Aとしては左を取るのが正解となり、このようにAが選ぶ以上、 Bも、最初の真ん中3つから残される値より小さい値を残すのは不可能となる。

AもBも、真ん中3つよりよい値を残したくても相手が阻止でき、その時に残る値は真ん中3つのみの時に残る値と一致する。

真ん中3数の大小関係の並びは複数あるが、調べれば、どのパターンもそうなる。

奇数個の場合 $O(1)$ で残る値が特定できるとわかったので、 偶数個の場合、先手は左右のどちらを取れば自分にとって望ましいかを2通りから選べる。

奇数長の数列は先に手を着けた方が不利。

奇数長の数列しか残っていない状況になると、 その時点での先手（奇数長の数列に手を付ける側）はその状況を覆せず、ずっと不利な状況のまま終わってしまう。

また、どちらがその可哀想な方になるかは、初期で偶数長の数列の個数によって決まっている。
偶数個ならA、奇数個ならBとなる。

偶数長の数列は先に手を着けた方が2通りから望む方を取れるので、率先して1回手をつけるべき。
よって、ゲームの流れは「偶数長の取り合い」→「奇数長の処理」となる。

どの偶数長の数列から手を付ければよいかは、利得の大きい方となる。

つまり、2通りの結果を $a,b$ として $\max(a,b)-\min(a,b)$ を計算する。 AもBもこの利得が大きい方から取っていくことになる。

Python3

k = int(input())
even_center = []
odd_center = []
two_que = []
ans = 0
for _ in range(k):
    n, *aaa = map(int, input().split())
    if n == 1:
        ans += aaa[0]
    elif n == 2:
        two_que.append(aaa)
    elif n % 2 == 0:
        even_center.append(aaa[(n - 3) // 2:(n - 3) // 2 + 4])
    else:
        odd_center.append(aaa[(n - 3) // 2:(n - 3) // 2 + 3])

even_count = len(two_que) + len(even_center)

if even_count % 2 == 0:
    can_take_center = lambda aaa: aaa[0] > aaa[1] < aaa[2]
else:
    can_take_center = lambda aaa: aaa[0] < aaa[1] > aaa[2]

for aaa in odd_center:
    if can_take_center(aaa):
        ans += aaa[1]
    else:
        ans += sorted(aaa)[1]

scramble = []
for aaa in two_que:
    a = min(aaa)
    b = max(aaa)
    scramble.append((a - b, a, b))
for aaa in even_center:
    if can_take_center(aaa[:3]):
        a = aaa[1]
    else:
        a = sorted(aaa[:3])[1]
    if can_take_center(aaa[1:]):
        b = aaa[2]
    else:
        b = sorted(aaa[1:])[1]
    if a > b:
        a, b = b, a
    scramble.append((a - b, a, b))

scramble.sort()

turn = 0
for diff, a, b in scramble:
    if turn == 0:
        ans += b
    else:
        ans += a
    turn ^= 1

print(ans)