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

AtCoder Regular Contest 107

A - Simple Math

• $A,B,C$ が与えられるので、$\displaystyle \sum_{a=1}^{A} \sum_{b=1}^{B} \sum_{c=1}^{C} abc$ を $\mod{998244353}$ で求めよ
• $1 \le A,B,C \le 10^9$

$a,b$ を固定して $c$ のループだけを考えると、以下のようになる。

• $(ab \times 1)+(ab \times 2)+...+(ab \times C)$

これは、$ab \times (1+2+...+C)$ とまとめられる。

すると、$\mathbf C=1+2+...+C$ として、$\mathbf C$ は定数なので外に出せて、最初の式は以下のようになる。

• $\displaystyle \mathbf C \sum_{a=1}^{A} \sum_{b=1}^{B} ab$

$b,a$ についても同様のことがいえるので、結局答えは $\mathbf A \mathbf B \mathbf C$ となる。

a, b, c = map(int, input().split())
ans = a * (a + 1) * b * (b + 1) * c * (c + 1) // 8 % 998244353
print(ans)

B - Quadruple

• 整数 $N,K$ が与えられる
• 以下の条件をともに満たす4つの整数の組 $(a,b,c,d)$ の個数を求めよ
• 条件
  • $1 \le a,b,c,d \le N$
  • $a+b−c−d=K$
• $1 \le N \le 10^5$
• $−2(N−1) \le K \le 2(N−1)$

$a+b$ と $c+d$ に分けて考える。

$a+b$ を固定すると、$c+d=(a+b)-K$ も決まる。また、$X=a+b$ の時に $a,b$ が取り得る組の個数も、値の範囲を考えるとすぐに計算できる。（同様に $c+d$ も）

なので、$a+b$ を全通り試す。

$a+b$ の取り得る範囲は、$c,d$ を移行すると $a+b=K+c+d$ となり、

• $2 \le a+b \le 2N$
• $2+K \le K+c+d \le 2N+K$

これが等しくなる（重なる）範囲のみ探せばよいので、$\max(2,2+K)～\min(2N,2N+K)$ となる。

Python3

n, k = map(int, input().split())
ans = 0
for ab in range(max(2, 2 + k), min(2 * n, 2 * n + k) + 1):
    max_a = min(n, ab - 1)
    min_a = max(1, ab - n)
    ab_pat = max_a - min_a + 1

    cd = ab - k
    max_c = min(n, cd - 1)
    min_c = max(1, cd - n)
    cd_pat = max_c - min_c + 1

    ans += ab_pat * cd_pat

print(ans)

C - Shuffle Permutation

• $N \times N$ の行列と整数 $K$ が与えられる
• 行列の $i$ 行目 $j$ 列目の要素は $a_{i,j}$ である。この行列の要素は $1～N^2$ が1つずつ出現する
• 以下の操作をどちらでも好きなだけ行える
  • 2つの行 $x,y$ について、同じ列の要素の和（$a_{x,j}+a_{y,j}$）が全ての列で $K$ 以下であれば、2行をスワップする
  • 2つの列 $x,y$ について、同じ行の要素の和（$a_{i,x}+a_{i,y}$）が全ての行で $K$ 以下であれば、2列をスワップする
• 操作後のあり得る行列の個数を $\mod{998244353}$ で求めよ
• $1 \le N \le 50$

まず、和が $K$ 以下とか関係なく、全ての行、列をスワップできるとすると、どういう行列を作れる？と考えると、

1 2 3      8 9 7
4 5 6  →  2 3 1
7 8 9      5 6 4

• 初期状態で同じ行の数字は（順番は変わっても、組は）変わらない
  • 1,2,3は互いに1つだけ別の行にいったりしない
• 初期状態で同じ列の数字は（順番は変わっても、組は）変わらない
  • 1,4,7は互いに1つだけ別の列にいったりしない

なので、行、列をそれぞれどういう順番で並べるかで $N! \times N!$ 通りが、制約の無いときに作れる行列となる。

また、この条件より、列をどれだけスワップしても、ある2行が入れ替え可能かどうかが変わることはないことがわかる。（逆も同じ）

さて、サンプル1は $K$ の制約無しでは $3! \times 3!=36$ 通りの行列が作れるところ、答えは12通りになっている。

3 2 7
4 8 9
1 6 5

これが何故かを観察すると、$K$ の制約があると '4 8 9' の行が動かせないから。
よって行方向では $2!$ 通りの並べ方しか無く、$2! \times 3!=12$ 通り、ということになる。

スワップできる行同士をUnion-Find等で連結していくと、同じ連結成分の行同士は（たとえ2つが直接スワップできなくても、適当な行を仲介すれば）好きに入れ替えられる。
その連結成分の最終的な並べ方は、連結成分のサイズを $S$ とすると $S!$ なので、これ毎に乗算していけばよい。

計算量は、2行を選ぶのが $O(N^2)$、選んでその2行が入れ替え可能か調べるのが $O(N)$、調べてUnion-Findで連結させるのはほぼ定数とみなせ¹⁾、全体で約 $O(N^3)$ となる。

Python3

import sys
import numpy as np


class UnionFind:
    def __init__(self, n):
        self.table = [-1] * n

    def _root(self, x):
        stack = []
        tbl = self.table
        while tbl[x] >= 0:
            stack.append(x)
            x = tbl[x]
        for y in stack:
            tbl[y] = x
        return x

    def find(self, x, y):
        return self._root(x) == self._root(y)

    def unite(self, x, y):
        r1 = self._root(x)
        r2 = self._root(y)
        if r1 == r2:
            return
        d1 = self.table[r1]
        d2 = self.table[r2]
        if d1 <= d2:
            self.table[r2] = r1
            self.table[r1] += d2
        else:
            self.table[r1] = r2
            self.table[r2] += d1

    def get_size(self, x):
        return -self.table[self._root(x)]


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, k, aaa, facts, MOD):
    uft = UnionFind(n)
    for i in range(n):
        for j in range(i + 1, n):
            if (aaa[i] + aaa[j] <= k).all():
                uft.unite(i, j)

    ans = 1

    for i in range(n):
        r = uft.table[i]
        if r < 0:
            ans = ans * facts[-r] % MOD

    return ans


n, k, *aaa = map(int, sys.stdin.buffer.read().split())
aaa = np.array(aaa).reshape((n, n))
MOD = 998244353

facts, _ = prepare(n, MOD)

rows = solve(n, k, aaa, facts, MOD)
cols = solve(n, k, aaa.T, facts, MOD)
print(rows * cols % MOD)

D - Number of Multisets

• 整数 $N,K$ が与えられる
• 以下の条件を全て満たす多重集合（同じ要素があってもいいが並び順は区別しない集合）の個数を $\mod{998244353}$ で求めよ
• 条件
  • 要素数は $N$
  • 要素の総和は $K$
  • 全ての要素は $1,\frac{1}{2},\frac{1}{4},...$ のように「2の累乗分の1」からなる
• $1 \le K \le N \le 3000$

制約から、なんとなく2次元DPっぽい。

• $DP[i][j]=$ 要素が $i$ 個、総和が $j$ の多重集合の個数

どう更新するかに、ひらめきが必要。

「$\frac{1}{4}$ を含む集合は $\frac{1}{8}$ 2つに分解して、総和を変えずに要素数を1増やせる」などと遷移しようとしても筋が悪い。
$DP[i][j]$ のうちいくつに $\frac{1}{4}$ が含まれいくつに含まれていないかという情報が必要になり、それを $1,\frac{1}{2},\frac{1}{4},...$ の全てではとても持たせていられない。

また、「1を $M$ 個の要素で作る場合の数 $D_M$ を求めておいて、$DP[i][j]+=DP[i][j-m] \times D_m$ を足し合わせる」みたいなことも、重複が発生しまくって、上手い除き方がわからない。

ここでは、以下の2つが同じことに気付けば見通しが立つ。

• $N$ 要素で総和が $2K$ の、$1,\frac{1}{2},\frac{1}{4},...$ からなる多重集合の個数
• $N$ 要素で総和が $K$ の、$\frac{1}{2},\frac{1}{4},\frac{1}{8},...$ からなる多重集合の個数

つまり、$DP[i][j]$ を求めたいとき、その中に“1”が存在するか否かで遷移元を変えて、

• 存在するとき
  • 1を1つ除いた場合の数をそのまま引き継ぐ
  • $DP[i-1][j-1]$
• 存在しないとき
  • 上記の言い換えで、要素数 $i$、総和 $2j$ の時の場合の数と等しい
  • $DP[i][2j]$
  • ただし問題設定から、$DP[i][i]=1$、また $i \gt j$ のとき $DP[i][j]=0$
• まとめて、
  • $DP[i][j]=DP[i-1][j-1]+DP[i][2j]$

$i$ の小さい方、$j$ の大きい方からDPテーブルを埋めていくことで、$O(N^2)$ で全て埋められる。

Python3

n, k = map(int, input().split())
MOD = 998244353

dp = [1]
for i in range(1, n + 1):
    ndp = [0] * (i + 1)

    for j in range(i, 0, -1):
        tmp = dp[j - 1]
        if j * 2 <= i:
            tmp += ndp[j * 2]
        ndp[j] = tmp % MOD

    dp = ndp

print(dp[k])

E - Mex Mat

• $N \times N$ のグリッドがあり、各要素は $0,1,2$ からなる
• 先頭行と先頭列は与えられる
• それ以外のマスの要素は、自身の1つ上と1つ左のマスの要素のMexとなる
  • Mexはその集合に含まれない最小の非負整数。詳細は問題ページ参照
• グリッド内の $0,1,2$ の要素の各個数を求めよ
• $1 \le N \le 5 \times 10^5$

こういうのは、問題文の操作通りにシミュレーションするコードを書いて実験する。
シミュレーションは簡単に書けるし、見た目上わかりやすい性質が表れてくれたら「何も無いところから性質を見つける」より「仮説を証明する」方が一般的には簡単。

すると、幾らもしないうちにナナメ↘に同じ数字が並び始めることが分かる。

2 1 0 0 1 1 2 1 1 1
2 0 1 2 0 2 0 2 0 2
2 1 0 1 2 0 1 0 1 0
0 2 1 0 1 2 0 1 0 1
0 1 0 1 0 1 2 0 1 0
0 2 1 0 1 0 1 2 0 1
1 0 2 1 0 1 0 1 2 0
2 1 0 2 1 0 1 0 1 2
2 0 1 0 2 1 0 1 0 1
2 1 0 1 0 2 1 0 1 0

ナナメに同じ数字が続くなら、実際にシミュレーションしなくても計算で個数は求められる。
必ずそうなるのか確かめる。

まず、所与である1行目・1列目を除いて、“0”の右下は必ず“0”となる。

0 A    A: 左に0があるので0にはならない
B C    B: 上に0があるので0にはならない
       C: 左も上も0ではないので0

また、（しばらく行が進んだ）隣接する3本のナナメ↘ラインには、必ず1つは“0”が含まれる。

↘↘
↘ A B        B,C がともに非0なら、Aは0でないと矛盾
   C A D      一度0になると先ほどの条件よりずっとAは0
     E A F

なので、“0”は、ナナメラインで1個おきか、2個おきに現れることとなる。
1個おきの場合は間の数字は1に収束し、2個おきは色々試すと 0 1 2 0 または 0 2 1 0 に収束する。

収束までにかかる回数はちょっと厳密に分からないが……

• 1行目は同じ数字が連続しうるので、「“0” 以外が隣接3本のナナメラインで続かない」性質が2行目の時点では現れないことがある
• 2行目は同じ数字が連続しないので、3行目ではその性質は保証される
• その時、“0”の間の数字が収束する値とは異なる場合があるが、4行目、5行目と繰り返す内にほぼすぐ収束する（はず）

よって、念のためでも行・列それぞれで10回程度シミュレーションすればまず収束し、個数を求められる。

収束するまでのシミュレーションは1行 or 1列あたり $O(N)$ でできる。

Python3

import sys
from collections import Counter


def mex(a, b):
    if a > 0 and b > 0:
        return 0
    elif a == 1 or b == 1:
        return 2
    else:
        return 1


def solve(n, cols, rows):
    cnt = [0, 0, 0]
    for c, v in Counter(cols).items():
        cnt[c] += v
    for c, v in Counter(rows[1:]).items():
        cnt[c] += v

    if n == 1:
        return cnt
    if n == 2:
        cnt[mex(cols[1], rows[1])] += 1
        return cnt

    cols_origin = cols.copy()
    simulate = min(n - 1, 10)

    for i in range(simulate):
        new_cols = [0] * n
        a = new_cols[0] = rows[i + 1]
        for j in range(1, n):
            a = new_cols[j] = mex(a, cols[j])
            cnt[a] += 1
        cols = new_cols

    for j in range(simulate):
        new_rows = [0] * n
        a = new_rows[0] = cols_origin[j + 1]
        for i in range(1, n):
            a = new_rows[i] = mex(a, rows[i])
            if i > simulate:
                cnt[a] += 1
        rows = new_rows

    for j, a in enumerate(cols[simulate:], start=simulate):
        cnt[a] += n - j - 1

    for i, a in enumerate(rows[simulate + 1:], start=simulate + 1):
        cnt[a] += n - i - 1

    return cnt


n = int(input())
cols = list(map(int, input().split()))
rows = [cols[0]] + list(map(int, sys.stdin.read().split()))

print(*solve(n, cols, rows))

D問題が苦手すぎたのが逆に功を奏してE問題に取りかかれたのがよかったかも。

F - Sum of Abs