AtCoder Regular Contest 190 (Div. 1) A,B,C,D問題メモ

AtCoder Regular Contest 190 (Div. 1)

A - Inside or Outside

問題文

整数列 $x = (x_1, \ldots, x_N)$ があり，$x_1=\cdots=x_N=0$ で初期化されています．
あなたはこの整数列について，$M$ 回の操作を行います．$i$ 回目の操作では，$1\leq L_i\leq R_i\leq N$ を満たす整数の組 $(L_i, R_i)$ が与えられるので，以下の $3$ つのうちちょうど $1$ つを行います．
- 操作 $0$：何もしない．この操作にはコストが $0$ かかる．
- 操作 $1$：$1\leq j\leq N$ を満たす各整数 $j$ に対して，$L_i\leq j\leq R_i$ を満たすならば $x_j=1$ と定める．この操作にはコストが $1$ かかる．
- 操作 $2$：$1\leq j\leq N$ を満たす各整数 $j$ に対して，$L_i\leq j\leq R_i$ を満たさないならば $x_j=1$ と定める．この操作にはコストが $1$ かかる．
あなたの目標は，最終的に $x_1=\cdots=x_N=1$ が成り立つようにすることです．この目標が達成できるか否かを判定してください．目標が達成可能な場合には，そのような方法のうち操作にかかるコストの総和が最小となるものをひとつ答えてください．

制約

$1\leq N\leq 1000000$
$1\leq M\leq 200000$
$1\leq L_i\leq R_i\leq N$
入力される値はすべて整数

解法

イメージとしては、$M$ 個の区間からなるべく少ない区間を選んで「区間内を塗る」か「区間外を塗る」かを行い、全てが塗られた状態にしてください、という問題。

丁寧場合分け。
達成できる場合は、3回以内で必ず達成可能。

答えが1の場合

これは明らかで、$(L_i,R_i)=(1,N)$ なる $i$ があったらOK。

$M=1$ で、答えが1の場合に当てはまらなければ、達成不可。

答えが2の場合

大きく分けて3通りの方法がある。

以下に置ける、区間の内外どちらを採用するかの記号の意味
      L    R
| - - |====| - - - -|    L～R の内を"1"にする
| = = |----| = = = =|    L～R の外を"1"にする

① 2つとも内を採用
|=========| - - - - |
| - - - |===========|

② 片方は内、片方は外を採用
②-a 左は内、右は外を採用
|=========| - - - - |
| = |---| = = = = = |
②-b 左は外、右は内を採用
| = = = = = |---| = |
| - - - |===========|
②-c 外採用の間を内採用の区間が埋める
| = = = |---| = = = |
| - - |=======| - - |

③ 2つとも外を採用。
| = = = = = |---| = |
| = |---| = = = = = |

①は、「$L_i=1$ の区間の内、$R_i$ 最大のもの」と「$R_j=N$ の区間の内、$L_j$ 最小のもの」を（存在すれば）比べればよい。

②は、どのケースを取っても、一方が一方に内包されていればよい。$(L_i,-R_i)$ でソートして確認できる。

③は、「$L_i$ 最大のもの」と「$R_j$ 最小のもの」を比べればよい。

$M = 2$ で、答えが1,2の場合に当てはまらなければ、達成不可。

答えが3の場合

「$L_i$ 最大のもの」と「$R_j$ 最小のもの」と「あと何でもいいのでもう1つ」で確実にできる。

      L       R
| = = |---------| = |  Li最大
| = |---------| = = |  Rj最小

|- - |=========| - -|  Li最大でもRj最小でもない区間は、
                       上の2つの間を必ず埋められる。

Python3

from operator import itemgetter


def solve(n, m, ops):
    ans = [0] * m

    ops = [(i, l - 1, r) for i, (l, r) in enumerate(ops)]
    op1_left = []  # l=1 の中で最大 r を持つものの (番号, l, r)
    op1_right = []  # r=n の中で最小 l を持つものの (番号, l, r)
    op2_left = []  # l!=1 の中で最大の l をもつものの (番号, l, r)
    op2_right = []  # r!=n の中で最小の r をもつものの (番号, l, r)
    for i, l, r in ops:
        if l == 0:
            op1_left.append((i, l, r))
        else:
            op2_left.append((i, l, r))
        if r == n:
            op1_right.append((i, l, r))
        else:
            op2_right.append((i, l, r))

    op1_left.sort(key=itemgetter(2), reverse=True)
    op1_right.sort(key=itemgetter(1))
    op2_left.sort(key=itemgetter(1), reverse=True)
    op2_right.sort(key=itemgetter(2))

    # print(op1_left)
    # print(op1_right)
    # print(op2_left)
    # print(op2_right)

    if op1_left and op1_left[0][2] == n:
        ans[op1_left[0][0]] = 1
        return 1, ans

    if m == 1:
        return -1, None

    if op1_left and op1_right:
        if op1_left[0][2] >= op1_right[0][1]:
            ans[op1_left[0][0]] = 1
            ans[op1_right[0][0]] = 1
            return 2, ans

    if op1_left and op2_right:
        for i, j in ((0, 0), (0, 1), (1, 0)):
            if len(op1_left) <= i or len(op2_right) <= j:
                continue
            if op1_left[i][0] == op2_right[j][0]:
                continue
            if op1_left[i][2] >= op2_right[j][2]:
                ans[op1_left[i][0]] = 1
                ans[op2_right[j][0]] = 2
                return 2, ans

    if op2_left and op1_right:
        for i, j in ((0, 0), (0, 1), (1, 0)):
            if len(op2_left) <= i or len(op1_right) <= j:
                continue
            if op2_left[i][0] == op1_right[j][0]:
                continue
            if op2_left[i][1] >= op1_right[j][1]:
                ans[op2_left[i][0]] = 2
                ans[op1_right[j][0]] = 1
                return 2, ans

    if op2_left[0][1] >= op2_right[0][2]:
        ans[op2_left[0][0]] = 2
        ans[op2_right[0][0]] = 2
        return 2, ans

    ops.sort(key=lambda a: (a[1], -a[2]))
    max_r = (0, 0)
    for i, l, r in ops:
        if max_r[0] >= r:
            ans[i] = 2
            ans[max_r[1]] = 1
            return 2, ans
        max_r = (r, i)

    if m == 2:
        return -1, None

    li = op2_left[0][0]
    ri = op2_right[0][0]
    ans[li] = 2
    ans[ri] = 2
    for i in range(m):
        if i != li and i != ri:
            ans[i] = 1
            break
    return 3, ans


n, m = map(int, input().split())
ops = [tuple(map(int, input().split())) for _ in range(m)]
cost, ans = solve(n, m, ops)
if cost == -1:
    print(-1)
else:
    print(cost)
    print(*ans)

B - L Partition

問題文

$N\times N$ のマス目があります．上から $i$ 行目，左から $j$ 列目のマスを $(i,j)$ で表します．
$K=1,2,\ldots,N$ に対して，レベル $K$ の L 型とは，以下のような形、及びこれを90,180,270度回転させた形を指します。
- ```
レベル1    レベル2    レベル3
             □       □
  □         □□     □
                      □□□
```

マス $(a,b)$ および $Q$ 個のクエリ $k_1, \ldots, k_Q$ が与えられます．
各 $i$ に対して，マス目全体をレベル $1, \ldots, N$ それぞれひとつずつの L 型に分割する方法であって，マス $(a,b)$ がレベル $k_i$ の L 型に含まれるようなものの個数を $998244353$ で割った余りを答えてください．

制約

$1\leq N\leq 10^7$
$1\leq a\leq N$
$1\leq b\leq N$
$1\leq Q\leq \min\lbrace N, 200000\rbrace$
$1\leq k_1 \lt \cdots \lt k_Q \leq N$
入力される値はすべて整数である

解法

レベル $k$ のL型の位置を固定した時の数え上げ

大きいレベルのL型から埋めて行くことを考えると、

レベル $N$ は、盤面をフルに使って4通りの転回形のどれかの形でしか入れられない
- →残る盤面は（位置こそ異なれど）$N-1 \times N-1$
レベル $N-1$ は、残った盤面をフルに使って4通りの転回形のどれかの形でしか入れられない
- →残る盤面は $N-2 \times N-2$
…

なので、$(a,b)$ に置く種類を無視すると、あり得る盤面は全部で $4^{N-1}$ 通りとなる。
レベル1のL型は転回しても同じなことに注意。

ここから、なんかタングラムのように複雑にL型が絡み合った配置にはなりようが無く、レベル $k-1$ 以下のL型は、全てレベル $k$ のL型のへこんだ部分に収まるしか無いということも分かる。

逆からの発想として、レベル $k$ のL型を置く位置を決めてしまった場合の、他のL型の配置方法の個数を考える。
（実装上、マス目のindexを0から始めるとする）

   0 1 2 3 4 5 6 7
0 □□□□□□□□    (2,3)～(5,6) の矩形に レベル4 のL型を置く
1 □□□□□□□□
2 □□□■■■■□
3 □□□○○○■□
4 □□□○○○■□
5 □□□○○○■□
6 □□□□□□□□
7 □□□□□□□□

この時、残る $3 \times 3$ の矩形（上記○印）をレベル3以下のL型で埋めるのは、先ほどの考え方で $4^{k-2}$ 通り。

矩形の外側をレベル5以上のL型で埋めるのは、$\binom{4}{2} \times \binom{4}{3}$ 通りとなる。
レベル5,6,… の順にキャベツの葉のように覆い被せながら配置を決めていく様子をイメージする。
矩形の外側には「縦に見ると上に2個、下に2個の空行」「横に見ると左に3個、右に1個の空列」が残っていて、 1つのL型で、1行1列を埋めることができる。
「上、左の組合せなら┌ 型」のように、転回すればどの ( {上,下}, {左, 右} ) の組にも対応できるので、縦と横は独立に決めてもよい。よって、$\binom{4}{2} \times \binom{4}{3}$ で求められる。

一般化すると、左上が $(i,j)$ である $k \times k$ の矩形の中にレベル $k$ のL字を向きを決めて置いたとすると、

$k=1$ の時、$\binom{N-k}{i} \times \binom{N-k}{j}$
$k \ge 2$ の時、$4^{k-2} \times \binom{N-k}{i} \times \binom{N-k}{j}$

通りの他のレベルの配置があり得ることになる。

よって、「$(a,b)$ にレベル $k$ のL型が置かれる」ような各場合に付き、上記を足し合わせれば求められる。

しかし、それでは1クエリ $O(N)$ かかり、全体で $O(NQ)$ でTLEしてしまう。

二項係数の累積和

$k=1$ の時は上記で問題なく求められるため、$k \ge 2$ とする。

レベル $k$ のL型を $(a,b)$ に置こうと思うと、 $k \times k$ の矩形の上下左右のいずれかの端を $(a,b)$ に合わせつつ、矩形全体が $N \times N$ に収まっていないといけない。

$k \times k$ の矩形の上、下、左、右それぞれの端を $(a,b)$ に合わせられるか確認し、合わせられるならその場合の答えを合計する。
ただし、$(a,b)$ を矩形の角に合わせられる場合、2回重複して数えられているので引く。
これで $k$ に対する答えが求められる。

たとえば右端を $(a,b)$ にあわせるためには、

   0 1 2 3 4 5 6 7       0 1 2 3 4 5 6 7    (a,b)=(2,6)  k=4
0 □□□■■■■□    0 □□□□□□□□
1 □□□□□□■□    1 □□□□□□□□    bを右端に合わせるとき、
2 □□□□□□●□    2 □□□■■■●□    左端の列は b-k+1 より、
3 □□□□□□■□    3 □□□□□□■□    これが >= 0 なら右端に合わせられる。
4 □□□□□□□□    4 □□□□□□■□
5 □□□□□□□□    5 □□□□□□■□    u = max(0, a-k+1) から
6 □□□□□□□□    6 □□□□□□□□    d = min(N-k, a)   までの i に対し
7 □□□□□□□□    7 □□□□□□□□    (i, b-k+1) を左上とした矩形が合わせられる。

この時、（上下ひっくり返したL型もあるので2倍して）矩形の右端に $(a,b)$ を合わせてレベル $k$ のL型を置き、それに合わせて他のL型も配置する方法は以下の数だけある。

$\displaystyle \sum_{i=u}^{d} 2 \times 4^{k-2} \times \binom{N-k}{i} \times \binom{N-k}{b-k+1}$
$\displaystyle = 2 \times 4^{k-2} \times \binom{N-k}{b-k+1} \times \sum_{i=u}^{d} \binom{N-k}{i}$

Σの外は $k$ と定数のみに依存するので、Σの中、パスカルの三角形の $N-k$ 段目の、左から $u～d$ 番目の項の合計が高速に分かれば良い。
これを求めるのはどうしても $O(N)$ かかり、普通にやると $O(NQ)$ のまま計算量は減ってない。

だが実は、クエリを通して様々な $k$ で必要になるこの区間和の区間 $[u,d]$ は、段毎に規則的な形となっている。

0        ●
1       ●●       ●: 区間和として各段で必要となる要素
2      ●●●
3     ●●●○     ←a+1 段目から右端が欠け始める
4    ●●●○○
5   ●●●○○○
6  ○●●○○○○  ←N-a 段目から左端が欠け始める
7 ○○●○○○○○

パスカルの三角形の各要素は下段に2回ずつ波及するので、区間和が大きく被っているとほぼ「$i$ 段目の区間和 = $i-1$ 段目の区間和 $\times 2$」となる。
端の要素だけが2回でなく1回だけ加算されるため、そこだけ調整すればよくなる。

類題: ABC235 G

よって、各段で必要となる区間和を $O(N)$ でまとめて事前計算できる。

今回は右端を合わせたが、上下左端を合わせる場合も同様に計算し、後は角で重複する分を引けば、$k$ に対する答えが求められる。

左右端をあわせる場合は上記のように $a$ を基準とした区間和が必要になるが、上下端を合わせる場合はまた別途 $b$ を基準とした区間和が必要になるので、事前計算は2通り必要になる。

Python3

import os
import sys

import numpy as np


def solve(inp):
    def mod_pow(x, a, MOD):
        ret = 1
        cur = x
        while a > 0:
            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

    MOD = 998244353
    n, a, b, q = inp[:4]
    kkk = inp[4:]
    a -= 1
    b -= 1
    facts, finvs = prepare_factorials(n, MOD)

    def ncr(n, k):
        return facts[n] * finvs[k] % MOD * finvs[n - k] % MOD

    def calc_cum_binom(a):
        cum_binom = np.ones(n, np.int64)
        for i in range(1, n):
            exclude = 0
            if i > a:
                exclude += ncr(i - 1, a)
            if i > n - a - 1:
                exclude += ncr(i - 1, n - a - 1)
            cum_binom[i] = (2 * cum_binom[i - 1] - exclude) % MOD
        return cum_binom

    cum_binom_a = calc_cum_binom(a)
    cum_binom_b = calc_cum_binom(b)
    # print('a', a, 'cum_binom_a', cum_binom_a)
    # print('b', b, 'cum_binom_b', cum_binom_b)
    pow4 = np.ones(n, np.int64)
    for i in range(1, n):
        pow4[i] = pow4[i - 1] * 4 % MOD

    result = np.zeros(q, np.int64)
    for qi in range(q):
        k = kkk[qi]
        if k == 1:
            result[qi] = cum_binom_a[n - 1] * cum_binom_b[n - 1] % MOD
            continue
        ans = 0

        coef_x = 0
        if b - k + 1 >= 0:
            coef_x += 2 * ncr(n - k, b - k + 1)
        if b + k <= n:
            coef_x += 2 * ncr(n - k, b)
        if coef_x > 0:
            ans += coef_x % MOD * cum_binom_a[n - k]
            ans %= MOD

        coef_y = 0
        if a - k + 1 >= 0:
            coef_y += 2 * ncr(n - k, a - k + 1)
        if a + k <= n:
            coef_y += 2 * ncr(n - k, a)
        if coef_y > 0:
            ans += coef_y % MOD * cum_binom_b[n - k]
            ans %= MOD

        if b - k + 1 >= 0 and a - k + 1 >= 0:
            ans -= ncr(n - k, b - k + 1) * ncr(n - k, a - k + 1) % MOD
        if b - k + 1 >= 0 and a + k <= n:
            ans -= ncr(n - k, b - k + 1) * ncr(n - k, a) % MOD
        if b + k <= n and a - k + 1 >= 0:
            ans -= ncr(n - k, b) * ncr(n - k, a - k + 1) % MOD
        if b + k <= n and a + k <= n:
            ans -= ncr(n - k, b) * ncr(n - k, a) % MOD

        ans %= MOD
        ans *= pow4[k - 2]
        ans %= MOD

        result[qi] = ans

    return result


SIGNATURE = '(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)

inp = np.fromstring(sys.stdin.read(), dtype=np.int64, sep=' ')
ans = solve(inp)
print('\n'.join(map(str, ans)))

C - Basic Grid Problem with Updates

問題文

$H\times W$ のマス目があります．上から $h$ 行目，左から $w$ 列目のマスを $(h,w)$ で表します．$(h,w)$ には非負整数 $A_{h,w}$ が書かれています．
高橋君がはじめマス $(sh,sw)$ におり，これからマス目に $Q$ 回の変更を行います．$i$ 回目の変更は文字 $d_i$ （$d_i$ は L, R, U, D のいずれか）と非負整数 $a_i$ によって与えられ，これは高橋君が以下を行うことを意味します．
- $d_i$ の方向に $1$ マス移動する．つまり，$d_i$ が L ならば左，$d_i$ が R ならば右，$d_i$ が U ならば上，$d_i$ が D ならば下方向に $1$ マス移動する．その後，移動先のマスを $(h,w)$ とするとき $A_{h,w}$ を $a_i$ に変更する．
- ただし，各変更において，高橋君が $d_i$ 方向に $1$ マス移動することが可能であることが保証されます．

マス目の変更のたびに，以下の問題の答を出力してください．
- マスの列 $P=((h_1,w_1),\ldots,(h_{M},w_{M}))$ が以下の条件をすべて満たすとき，$P$ はパスであるということにします．
  - $(h_1,w_1)=(1,1)$, $(h_{M},w_{M})=(H,W)$, $M=H+W-1$.
  - $1\leq i \leq M-1$ を満たす任意の $i$ に対して，$(h_{i+1},w_{i+1}) = (h_i+1,w_i)$ または $(h_{i+1},w_{i+1}) = (h_i,w_i+1)$ が成り立つ．
- パスは $\binom{H+W-2}{H-1}$ 個あります．パス $P=((h_1,w_1),\ldots,(h_{M},w_{M}))$ に対して $f(P) = \prod_{1\leq i\leq M}A_{h_i,w_i}$ と定めます．
- すべてのパス $P$ に対する $f(P)$ の総和を $998244353$ で割った余りを出力してください．

制約

$2\leq H, W\leq 200000$
$HW\leq 200000$
$0\leq A_{h,w} \lt 998244353$
$1\leq Q\leq 200000$
$1\leq sh\leq H$, $1\leq sw\leq W$
$0\leq a_i \lt 998244353$
$H, W, A_{h,w}, Q, sh, sw, a_i$ は整数
$d_i$ は L, R, U, D のいずれか
各変更において，高橋君が $d_i$ 方向に $1$ マス移動することが可能である

解法

解法はまぁまぁ自然な発想で導けるが、上手く実装するのが難しい。

問題の定義では「パス」は $(1,1)$ から $(H,W)$ に至るもののみを指すが、以下ではその制約は外し、任意の始点・終点のものを表すとする。

$\mathrm{DP}_{F}[i,j]:=(1,1)$ から $(i,j)$ に至る全パスにおける $f(P)$ の総和
$\mathrm{DP}_{B}[i,j]:=(i,j)$ から $(H,W)$ に至る全パスにおける $f(P)$ の総和

これを計算しておく。

このまま変更が無い場合の答えは（もちろん $\mathrm{DP}_F[H,W]$ を見れば一発なのだが）、ある行 $y$ に対して

$\displaystyle \sum_{j=1}^{W}\mathrm{DP}_F[y,j] \mathrm{DP}_B[y+1,j]$

としても求められる。「$y→y+1$ 行目の移動を $j$ 列目で行った場合」をそれぞれ足していると考えられる。

あるマスを変更したとき、一般的にはそこから先全てのDP配列の再計算が必要となってしまうが、今回は「変更するマスが隣接している」という特徴がある。

よって、$A_{i,j}$ が変更されるたびにそれぞれ1行ずつだけ再計算し、

高橋君の今いる位置を $(y,x)$ として
- $\mathrm{DP}_F$ は、$i=1～y$ までは正しい（$y+1$ 以降は未再計算）
- $\mathrm{DP}_B$ は、$i=y+1～H$ までは正しい（$y$ 以前は未再計算）

という状態を保っておくと、2つを掛け合わせることで、クエリ当たり $O(W)$ で答えられるようになる。

$H \lt W$ の場合は転置して考えればよいため、計算量は $O(HW+\min(H,W)Q)$ となる。

Python3

import os
import sys

import numpy as np


def solve(inp):
    h, w = inp[:2]
    aaa = np.ascontiguousarray(inp[2:2 + h * w]).reshape(h, w)
    q, si, sj = inp[2 + h * w:5 + h * w]
    queries = np.ascontiguousarray(inp[5 + h * w:]).reshape(q, 2)

    if h < w:
        h, w = w, h
        aaa = aaa.T
        si, sj = sj, si
        queries[:, 0] ^= 0b10

    bbb = np.zeros((h + 2, w + 2), np.int64)
    bbb[1:h + 1, 1:w + 1] = aaa
    bbb[0, 1] = bbb[-1, -2] = 1
    aaa = bbb

    MOD = 998244353
    fwd = aaa.copy()
    bwd = aaa.copy()
    for i in range(1, h + 1):
        for j in range(1, w + 1):
            fwd[i, j] *= fwd[i - 1, j] + fwd[i, j - 1]
            fwd[i, j] %= MOD
    for i in range(h, 0, -1):
        for j in range(w, 0, -1):
            bwd[i, j] *= bwd[i + 1, j] + bwd[i, j + 1]
            bwd[i, j] %= MOD

    ans = np.zeros(q, np.int64)
    i = si
    j = sj
    for qi in range(q):
        d, a = queries[qi]
        if d == 0:
            j -= 1
        elif d == 1:
            j += 1
        elif d == 2:
            i -= 1
        else:
            i += 1
        aaa[i, j] = a
        for l in range(1, w + 1):
            fwd[i, l] = aaa[i, l] * (fwd[i - 1, l] + fwd[i, l - 1]) % MOD
        for l in range(w, 0, -1):
            bwd[i, l] = aaa[i, l] * (bwd[i + 1, l] + bwd[i, l + 1]) % MOD
        if i == h:
            ans[qi] = fwd[h, w]
        else:
            tmp = 0
            for l in range(1, w + 1):
                tmp += fwd[i, l] * bwd[i + 1, l]
                tmp %= MOD
            ans[qi] = tmp

    return ans


SIGNATURE = '(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)

inp = np.fromstring(sys.stdin.read().translate(str.maketrans('LRUD', '0123')), dtype=np.int64, sep=' ')
ans = solve(inp)
print('\n'.join(map(str, ans)))

D - Matrix Pow Sum

問題文

素数 $p$ と $N\times N$ 行列 $A = (A_{i,j})$ ($1\leq i,j\leq N$) が与えられます．
$A$ の成分は $0$ 以上 $p-1$ 以下の整数です．
$A$ の成分のうち，$0$ を全て $1$ 以上 $p-1$ 以下の整数に置き換えて得られる行列 $B$ は，$A$ に含まれる $0$ の個数を $K$ とすると $(p-1)^K$ 個あります．
考えられる $B$ 全てに対する $B^p$ の総和の各成分を $p$ で割った余りを求めてください．

制約

$1\leq N\leq 100$
$p$ は $1\leq p\leq 10^9$ を満たす素数
$0\leq A_{i,j}\leq p-1$
入力される値はすべて整数である

解法

超・勘が鋭ければ気付けなくもないが、基本的には知らないとつらい。

重要なのは、以下の事実である。

3以上の素数 $p$ と、ある整数 $k$ に対し、$\displaystyle f(p,k)=\sum_{x=1}^{p-1}x^k$ を考える。

このとき、$f(p,k) \not\equiv 0 \mod{p}$ となるのは、$k$ が $p-1$ の倍数の時に限られ、その時 $f(p,k) \equiv -1$ である。

略証

まず、$p=2$ のとき、考えられる行列は1通りしかないので、素直に計算できる。以下、$p \ge 3$ とする。

$A_{i,j}=0$ である箇所をそれぞれ別の変数 $x_{i,j}$ に置き換え、変数含みの多項式行列として累乗を計算することをイメージする。

$$ \begin{pmatrix} x_{1,1} & 1 \\ x_{2,1} & 2 \end{pmatrix}^3= \begin{pmatrix} x_{1,1}^3+2x_{1,1}x_{2,1}+2x_{2,1} & x_{1,1}^2+2x_{1,1}+x_{2,1}+4 \\ x_{1,1}^2x_{2,1}+x_{2,1}^2+2x_{1,1}x_{2,1}+4x_{2,1} & x_{1,1}x_{2,1}+4x_{2,1}+8 \end{pmatrix} $$

ここで、各 $x_{i,j}$ を $1～p-1$ に置き換えた結果の総和を取るというのは、各 $x_{i,j}^k$ を $\displaystyle \sum_{y=1}^{p-1}y^k$ に置き換えることに相当する。

ここでさっきの定理が活き、ほとんどの項は mod p で 0 となる。
全ての変数の次数が $p-1$ の倍数となる項だけを計算すればよい。

「全ての変数の次数が $p-1$ の倍数となる項」とは、以下の2つに限られる。

定数項
ある単一の変数 $x$ で $(係数) \times x^{p-1}$ として表される項

2つ以上の変数を含んだ項はあり得ない。
∵一次多項式行列を $k$ 乗した結果、各項の次数（$x^py^qz^r$ などにおける $p+q+r$）が $k$ を超えることはない。

定数項は普通に行列累乗で計算する。
入力をそのまま（変数の箇所は $0$ としたまま）$p$ 乗すると、変数を含むはずの項は $0$ で消されるので、定数項のみの結果が得られる。

単一変数の項は、ある程度、出現する場所は限られるはずとにらんで実験すると、以下の法則がある。

対角線上にある変数について、それと同じ行または列で、$A_{i,j} \neq 0$ である箇所に、係数 $A_{i,j}$ で出現
$p=3$ の場合に限り、対角線上にない変数 $x_{i,j}$について、$A_{j,i} \neq 0$ である場合、$(i,j)$ に係数 $A_{j,i}$ で出現

A             A^k
x 1 2 0       - 1x 2x  -    ※A^k の結果の "x" は、実際は x^(k-1)
3 - - -  →  3x  -  -  -
0 - - -       -  -  -  -
8 - - -      8x  -  -  -

A             A^3
- - x -       -  - 2x  -
- - - y  →   -  -  - 3y
2 - - -       -  -  -  -
- 3 - -       -  -  -  -

よって、この位置のみ計算することで、$O(N^3)$ で影響を全て求められる。

最後に、出現する変数の個数を $z$ として、定数項の場合は $(p-1)^z$、単一変数の場合は $(p-1)^{z-1}$ だけ他の変数の決め方があるので、倍加した上で合計すると答えとなる。（単一変数の係数は $-1$ で寄与する点に注意）

定数項の行列累乗がボトルネックとなり、$O(N^3 \log{p})$ で求められる。

$x^{p-1}$ の出現位置にわかりやすい法則があることに気付く前は、「行列の各要素を、変数含みの式を表す辞書型とし、ダブリングで行列累乗する際、$A^k$ に対して、単一変数で次数が $k$ と $k-1$ の項だけ持つ」を検討していたが、TLEが取れなかった。
累乗を重ねると各辞書の要素数は高々2～3個となるはずなのでギリ間に合うかと思ったが、やはり $O(N^3 \log{p})$ に追加の定数倍はつらかったか。

Python3

import os
import sys

import numpy as np


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

    def matrix_mul(n, a, b, MOD):
        result = np.zeros((n, n), np.int64)
        for i in range(n):
            for j in range(n):
                for k in range(n):
                    result[i, j] += a[i, k] * b[k, j]
                    result[i, j] %= MOD
        return result

    n, p = inp[:2]
    mat = np.ascontiguousarray(inp[2:]).reshape(n, n)

    if p == 2:
        return np.full((n, n), n % 2, np.int64)

    matrix_const = mat.copy()
    result_const = np.eye(n, n, dtype=np.int64)
    bit = p
    while bit > 0:
        if bit & 1:
            result_const = matrix_mul(n, result_const, matrix_const, p)
        matrix_const = matrix_mul(n, matrix_const, matrix_const, p)
        bit >>= 1

    result_var = np.zeros((n, n), np.int64)
    zero = 0
    for i in range(n):
        for j in range(n):
            if mat[i, j] != 0:
                continue
            zero += 1
            if i == j:
                for k in range(n):
                    if mat[i, k] != 0:
                        result_var[i, k] += mat[i, k]
                    if mat[k, j] != 0:
                        result_var[k, j] += mat[k, j]
            elif p == 3 and mat[j, i] != 0:
                result_var[i, j] += mat[j, i]
    result_var %= p

    result_const *= mod_pow(p - 1, zero, p)
    result_const %= p
    result_var *= -mod_pow(p - 1, zero - 1, p)
    result_var %= p

    result = (result_const + result_var) % p
    return result


SIGNATURE = '(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)

inp = np.fromstring(sys.stdin.read(), dtype=np.int64, sep=' ')
ans = solve(inp)
for row in ans:
    print(*row)

目次

AtCoder Regular Contest 190 (Div. 1) A,B,C,D問題メモ

A - Inside or Outside

問題文

制約

解法

答えが1の場合

答えが2の場合

答えが3の場合

B - L Partition

問題文

制約

解法

レベル $k$ のL型の位置を固定した時の数え上げ

二項係数の累積和

C - Basic Grid Problem with Updates

問題文

制約

解法

D - Matrix Pow Sum

問題文

制約

解法