サントリープログラミングコンテスト2024（AtCoder Beginner Contest 357）G問題メモ

サントリープログラミングコンテスト2024（AtCoder Beginner Contest 357）

G - Stair-like Grid

問題文

$N$ 行の特殊な形をしたグリッドがあります。($N$ は偶数)
上から $i$ 行目にはマスが左端から $\left \lceil \frac{i}{2} \right \rceil \times 2$ マス並んでいます。
上から $i$ 行目、左から $j$ 列目のマスを $(i, j)$ と表します。

N = 6
□□
□□
□□□□
□□□□
□□□□□□
□□□□□□

各マスは空きマスと壁マスのいずれかです。壁マスは $M$ 個あり、$i$ 個目の壁マスは $(a_i, b_i)$ にあります。ただし $(1, 1)$ と $(N, N)$ は空きマスです。
$(1, 1)$ を出発して、右または下に隣り合う空きマスへの移動を繰り返して $(N, N)$ へ辿り着く方法は何通りありますか？答えを $998244353$ で割った余りを求めてください。

制約

$2 \leq N \leq 2.5 \times 10^5$
$N$ は偶数
$0 \leq M \leq 50$
$1 \leq a_i \leq N$
$1 \leq b_i \leq \left \lceil \frac{a_i}{2} \right \rceil \times 2$
$(a_i, b_i) \neq (1, 1)$ かつ $(a_i, b_i) \neq (N, N)$
$i \neq j$ ならば $(a_i, b_i) \neq (a_j, b_j)$

解法

普通のグリッドの場合

公式Editorialにもあるが、通常の $H \times W$ のグリッド中に（2乗が通る程度の）$M$ 個の壁マスがある、という問題であれば、EDPC Y - Grid 2 で既出。

全体から壁マスを通過する経路を引く。壁マスも通過できるとして、（確実に）踏んだ壁マス数により包除原理DPをおこなう。

$\mathrm{DP_{仮}}[i,k]=$ いま壁マス $i$ にいて、これまでに踏んだ壁マス数が $k$ 個以上の場合の経路数

さらに、「壁マス $i$ にいる状態から、次に壁マス $j$ に行く」という遷移は、これまでに踏んだ黒マス数 $k$ に依らず一定である。（共通して、$i→j$ への経路数がかかる）
こういう場合、個数すら省略して、偶奇（正にかかるか負にかかるか）だけにまとめてしまえる。

$\mathrm{DP}[i]=$ いま壁マス $i$ にいる場合の、（踏んだ黒マスが偶数の場合の経路数）-（奇数の場合の経路数）

計算しやすいグリッドにして、壁マスを追加

               ＼
               ○＼      ○: 追加した空きマス（始点・終点）
□□           □□＼    ■: 追加した壁マス
□□      →   □□■＼
□□□□       □□□□＼
□□□□       □□□□■＼
□□□□□□   □□□□□□＼
□□□□□□   □□□□□□○＼

$N+2$ の階段状のグリッドにして、$i=j$ のナナメラインのマスは踏んではいけないとする。
また、2段に1段、一番右のマスは踏んではいけないとする。
この上で、左上の○→右下の○の経路数は、元の問題と同じである。

こういうグリッドの $(i_1,j_1)→(i_2,j_2)$ の経路数は、カタラン数や鏡像法を使って、階乗の事前計算のもとで $O(1)$ で求められる。

2点間の経路数が高速に計算できるなら、EDPC Y と同等の手法が使える。。。
のだが、全ての壁マスを同等に扱うと $O((N+M)^2)$ かかるのでTLE。

本来の壁マスと異なり、追加した壁マス（「角マス」とする）同士の遷移は規則的になることを利用して、分割統治FFTをする。

ある角マス $i$ に対して $0～i-1$ の角マスのDPが分かっている場合、角マス間の遷移はカタラン数の偶数番目 $c_2,c_4,c_6,...$ を、$DP[i-1],DP[i-2],...$ と掛け合わせたものの合計となる。

○----- x132-,     DP[i-3] x 132
□□         |
□□■-- x14-|     DP[i-2] x  14
□□□□     |
□□□□■x2-|   + DP[i-1] x   2
□□□□□□ v   ---------------
□□□□□□■     DP[i]

こういう形のDPは、分割統治FFTで $O(N \log^2{N})$ で求められる。

解説 - AtCoder Beginner Contest 213

あとは、本来の壁マスの遷移を、分割統治の間に挟み込む。
角マスのDPを $\mathrm{DP}_1$、壁マスのDPを $\mathrm{DP}_2$ のように、分けて管理する。

「遷移元となるのは、値が確定してもう更新されないDP」（そうしないと、遷移元が更新されるたびに遷移先も更新が発生する）という点に注意すると、角マスの $\mathrm{DP}_1[i]$ が確定するのは、分割統治FFTで $[i,i+1)$ の処理に来たときである。
そのタイミングで、壁マスに関する遷移を計算する。

確定した角マス i からは、そこから右下の壁マスに遷移できる

○
□□ i
□□■        ×: DP1[i] が確定した際、i からの遷移を計算する壁マスの範囲
□□××
□□××■
□□××××
□□××××■

i からの遷移を計算し終えたら、i の下にある壁マスも（壁マス同士を除いては）確定する

○
□□ i
□□■        ×: DP1[i] から壁マスへの遷移を計算した後、
□□××          DP2[j] が確定する範囲（壁マス同士を除いて）
□□××■
□□××□□    →×の中で左上の壁マスから順次確定させていける
□□××□□■

「×内の壁マス→それより右下の壁マス・角マスの遷移」を左上から計算し、DPに足し込んでおく。

その後、分割統治FFTの続きを再開すれば、正しいDPの値となる。

壁マスにまつわる遷移は、$M$ 個それぞれで、計 $O(N+M)$ 個の他の黒マス {から/へ} の遷移を計算するため、$O(M(N+M))$ かかる。

あわせて、$O(N \log^2{N} + M(N+M))$ となる。

PythonだとTLEがやや厳しいが、畳み込み部分だけでもNumbaで高速化したりすると通る。

Python3

import os
import sys

import numba
import numpy as np

MOD = 998244353
SUM_E = np.array(
    [911660635, 509520358, 369330050, 332049552, 983190778, 123842337, 238493703, 975955924, 603855026, 856644456,
     131300601, 842657263, 730768835, 942482514, 806263778, 151565301, 510815449, 503497456, 743006876, 741047443,
     56250497, 867605899, 0, 0, 0, 0, 0, 0, 0, 0], np.int64)
SUM_IE = np.array(
    [86583718, 372528824, 373294451, 645684063, 112220581, 692852209, 155456985, 797128860, 90816748, 860285882,
     927414960, 354738543, 109331171, 293255632, 535113200, 308540755, 121186627, 608385704, 438932459, 359477183,
     824071951, 103369235, 0, 0, 0, 0, 0, 0, 0, 0], np.int64)


def convolution(aaa, bbb):
    def bit_length(n):
        x = 0
        while 1 << x < n:
            x += 1
        return x

    def bit_scan_forward(n):
        x = 0
        while n & 1 == 0:
            n >>= 1
            x += 1
        return x

    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 butterfly(aaa, mod, sum_e):
        n = aaa.size
        h = bit_length(n)

        for ph in range(1, h + 1):
            w = 1 << (ph - 1)
            p = 1 << (h - ph)
            now = 1
            for s in range(w):
                offset = s << (h - ph + 1)
                for i in range(p):
                    l = aaa[i + offset]
                    r = aaa[i + offset + p] * now % mod
                    aaa[i + offset] = (l + r) % mod
                    aaa[i + offset + p] = (l - r) % mod
                now = now * sum_e[bit_scan_forward(~s)] % mod

    def butterfly_inv(aaa, mod, sum_ie):
        n = aaa.size
        h = bit_length(n)

        for ph in range(h, 0, -1):
            w = 1 << (ph - 1)
            p = 1 << (h - ph)
            inow = 1
            for s in range(w):
                offset = s << (h - ph + 1)
                for i in range(p):
                    l = aaa[i + offset]
                    r = aaa[i + offset + p]
                    aaa[i + offset] = (l + r) % mod
                    aaa[i + offset + p] = ((l - r) * inow) % mod
                inow = inow * sum_ie[bit_scan_forward(~s)] % mod

    n = aaa.size
    m = bbb.size
    k = n + m - 1
    z = 1 << bit_length(k)
    raaa = np.zeros(z, np.int64)
    rbbb = np.zeros(z, np.int64)
    raaa[:n] = aaa
    rbbb[:m] = bbb
    butterfly(raaa, MOD, SUM_E)
    butterfly(rbbb, MOD, SUM_E)
    for i in range(z):
        raaa[i] = raaa[i] * rbbb[i] % MOD
    butterfly_inv(raaa, MOD, SUM_IE)
    iz = mod_pow(z, MOD - 2, MOD)
    for i in range(k):
        raaa[i] = raaa[i] * iz % MOD
    return raaa[:k]


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

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

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

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


def precompute_factorials(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


input = sys.stdin.buffer.readline
n, m = map(int, input().split())
facts, finvs = precompute_factorials(n * 2, MOD)
n2 = n // 2

# 偶数番目のカタラン数リスト
catalan_even = [facts[i * 2] * finvs[i] * finvs[i + 1] % MOD for i in range(0, n + 1, 2)]
catalan_even = np.array(catalan_even, dtype=np.int64)

dp1 = np.zeros(n2 + 1, dtype=np.int64)
dp2 = [0] * m
dp1[0] = 1

walls = []
for wi in range(m):
    i, j = map(int, input().split())
    i += 1
    j -= 1
    walls.append((j, i))
walls.sort()
wall_index = []
wi = 0
for i in range(n2 + 1):
    while wi < m and walls[wi][0] // 2 < i:
        wi += 1
    wall_index.append(wi)
wall_index.append(m)


def get_path_count(si, sj, ti, tj):
    # i=j のナナメライン上にあるマスを通らない経路数
    di1 = ti - si
    dj1 = tj - sj
    di2 = tj - si
    dj2 = ti - sj
    assert di1 >= 0 and dj1 >= 0
    res = facts[di1 + dj1] * finvs[di1] * finvs[dj1]
    if di2 >= 0 and dj2 >= 0:
        res -= facts[di2 + dj2] * finvs[di2] % MOD * finvs[dj2]
    res %= MOD
    return res


def divide_and_conquer(l, r):
    if l + 1 == r:
        # 角マス dp1[l] が確定したので、lより右下の壁マスの dp2 に遷移
        xi = l * 2 + 1
        xj = l * 2
        dp1l = dp1[l]
        for wi in range(wall_index[l], m):
            j, i = walls[wi]
            dp2[wi] -= dp1l * get_path_count(xi, xj, i, j)
            dp2[wi] %= MOD
        # l の下の壁マスについて、角マス→壁マスの遷移は確定したので、壁→壁 の遷移を計算
        for wi1 in range(wall_index[l], wall_index[l + 1]):
            j1, i1 = walls[wi1]
            dp2w = dp2[wi1]
            for wi2 in range(wi1 + 1, m):
                j2, i2 = walls[wi2]
                if i1 <= i2:
                    dp2[wi2] -= dp2w * get_path_count(i1, j1, i2, j2)
                    dp2[wi2] %= MOD
            # l の下の壁マスの dp2 が確定したので、壁マスより右下にある角マス dp1 に遷移
            for ci in range(i1 // 2, n2 + 1):
                xi = ci * 2 + 1
                xj = ci * 2
                dp1[ci] -= dp2w * get_path_count(i1, j1, xi, xj)
            dp1[i1 // 2:] %= MOD
        return

    mid = (l + r) // 2
    divide_and_conquer(l, mid)

    left = dp1[l:mid]
    right = catalan_even[:r - l]

    cnv = convolution(left, right)
    dp1[mid:r] -= cnv[mid - l:r - l]
    dp1[mid:r] %= MOD

    divide_and_conquer(mid, r)


divide_and_conquer(0, n2 + 1)

ans = -dp1[n2] % MOD
print(ans)

目次

サントリープログラミングコンテスト2024（AtCoder Beginner Contest 357）G問題メモ

G - Stair-like Grid

問題文

制約

解法

普通のグリッドの場合

計算しやすいグリッドにして、壁マスを追加