AtCoder Beginner Contest 265 F問題メモ

AtCoder Beginner Contest 265

実装重めハッピーセット。

F - Manhattan Cafe

問題

$N$ 次元空間に点 $P=(P_1,...,P_N)$ と点 $Q=(Q_1,...,Q_N)$ がある
$P$ からも $Q$ からもマンハッタン距離が $D$ 以下の格子点がいくつあるか、$\mod{998244353}$ で求めよ
$1 \le N \le 100$
$0 \le D \le 1000$
実行時間制限: $6$ sec.

解法

$O(ND^2)$ で頑張るDP。

$DP[i][j][k]=i$ 次元目までで考えて、$P$ からの各次元の距離の和が $j$、$Q$ からのそれが $k$ な格子点の個数

$DP[0][0][0]=1$ からはじめて、1次元目はどうなるかというと、

D=5    位置 -2     0        3     5
             |-----P--------Q-----|
Pからの距離  2  1  0  1  2  3  4  5
Qからの距離  5  4  3  2  1  0  1  2

j\k 0 1 2 3 4 5
---+-----------
 0 |      1
 1 |    1   1
 2 |  1       1
 3 |1
 4 |  1
 5 |    1

たとえば $(0,0)$ を遷移元とした場合、そこから $|P_i-Q_i|$ だけ右下に動いた位置に「／」方向に $1$ が並び、もしそれが遷移元と同じ行や列に来たら、そこで反射したような形になる。

1次元目は $DP[0][0][0]=1$ からしか遷移しないのでシンプルに $1$ が並ぶだけだが、以降の遷移は「DPの各マスから、相対的にこのような位置に遷移する」。

遷移先は直線上に並んでいて $O(D)$ あるので、 1次元分の遷移につき $D^2$ マスからこれを繰り返すと $O(ND^3)$ となり、これはさすがにTLE。

直線上に並んでいるということを利用して、累積和で更新する。

「↙」用の累積和と「↘」用の累積和用テーブルを用意して、

(0,0) からの遷移
↙                ↘
j\k 0 1 2 3 4 5   j\k 0 1 2 3 4 5
---+-----------   ---+-----------
 0 |      1        0 |
 1 |               1 |        1
 2 |               2 |
 3 |               3 |
 4 |               4 |  1
 5 |               5 |

こうして各マスからの遷移を直線の始終点だけ更新すればよいようにしておいて、全マスの遷移を処理した後に累積和を取って

↙                ↘
j\k 0 1 2 3 4 5   j\k 0 1 2 3 4 5
---+-----------   ---+-----------
 0 |      1        0 |
 1 |    1          1 |        1
 2 |  1            2 |          1
 3 |1              3 |
 4 |               4 |  1
 5 |               5 |    1

あわせれば

j\k 0 1 2 3 4 5
---+-----------
 0 |      1
 1 |    1   1
 2 |  1       1
 3 |1
 4 |  1
 5 |    1

1次元当たりの更新を $O(D^2)$ でおこなえ、全体で $O(ND^2)$ となる。

もらうDP

上記は配るDPで考えたが、もらうDPで考えることもできる。まぁ、考え方自体は変わらない。

j\k 0 1 2 3 4 5
---+-----------
 0 |      1
 1 |        1
 2 |          1
 3 |1       1
 4 |  1   1
 5 |    1     *

配るときとちょうど点対称のように、たとえば $(5,5)$ へ遷移するマスは、左上へ $|P_i-Q_i|$ だけ動いた「／」と、遷移先と同じ行・列での反射となる。

$DP[i-1]$ から $DP[i]$ へは「↘の累積和と↙の累積和」の状態で引き継いでやれば、 $DP[i]$ の各マスを $O(1)$ で計算でき、次に引き継ぐための「↘の累積和、↙の累積和」も $O(D^2)$ で計算できる。

Python3

n, d = map(int, input().split())
ppp = list(map(int, input().split()))
qqq = list(map(int, input().split()))
d2 = d * 2
MOD = 998244353

dp_acc_l = [[0] * (d + 2) for _ in range(d2 + 2)]  # ↙累積和[i+j][i]
dp_acc_r = [[0] * (d + 2) for _ in range(d2 + 2)]  # ↘累積和[i-j][i]
dp_acc_l[0][0] = 1
for i in range(d + 1):
    dp_acc_r[0][i] = 1

for pi in range(n):
    dp = [[0] * (d + 1) for _ in range(d + 1)]
    p = ppp[pi]
    q = qqq[pi]
    k = abs(p - q)
    for i in range(d + 1):
        for j in range(d + 1):
            if i + j < k:
                continue
            s = max(0, i - k)
            t = i - max(0, k - j)
            dp[i][j] = dp_acc_l[i + j - k][t] - dp_acc_l[i + j - k][s - 1]

            if i > k:
                dp[i][j] += dp_acc_r[i - j - k][i - k - 1]

            if i > 0 and j > k:
                dp[i][j] += dp_acc_r[i - j + k][i - 1]

    dp_acc_l = [[0] * (d + 2) for _ in range(d2 + 2)]  # ↙累積和[i+j][i]
    dp_acc_r = [[0] * (d + 2) for _ in range(d2 + 2)]  # ↘累積和[i-j][i]
    for i in range(d + 1):
        for j in range(d + 1):
            dp_acc_l[i + j][i] = (dp_acc_l[i + j][i - 1] + dp[i][j]) % MOD
            dp_acc_r[i - j][i] = (dp_acc_r[i - j][i - 1] + dp[i][j]) % MOD

ans = sum(sum(row) for row in dp) % MOD
print(ans)

目次

AtCoder Beginner Contest 265 F問題メモ

F - Manhattan Cafe

問題

解法

もらうDP