HHKBプログラミングコンテスト2022 Winter(AtCoder Beginner Contest 282) G問題メモ

HHKBプログラミングコンテスト2022 Winter(AtCoder Beginner Contest 282)

G - Similar Permutation

• $1～N$ の順列を2つ作ることを考える
  • $A=(A_1,...,A_N)$、$B=(B_1,...,B_N)$
• $i=1～N-1$ につき、「右隣の要素が自身より大きいか小さいか」が $A$ と $B$ で一致している箇所の個数を「類似度」とする
  • 数式的に言うと、$(A_{i+1}-A_i)(B_{i+1}-B_i) \gt 0$
• 類似度がちょうど $K$ になるような $A,B$ の組の個数を、与えられる素数 $P$ で割ったあまりで求めよ
• $2 \le N \le 100$

1回の遷移に $O(N^2)$ かかる4次元DPで、全体で $O(N^6)$ かかるところを、2次元累積和で $O(N^4)$ に高速化。

挿入DPと言われるらしいが、個人的なイメージでは挿入と言うよりダルマ落としのように抜いていく感覚。
挿入DPと検索して出てくる解説ページを見ても説明の仕方が様々で、これが本当に一緒の種類のDPなの？ となって、この言葉の指す範囲がちゃんと分かっていない。

先頭から順番に、順列の要素を $A,B$ 同時に決めていく。

$A_{i+1},B_{i+1}$ に置く値を決めて類似度がどうなるか（大小が一致して1増えるか、変わらないか）を判断するには、 $A_i,B_i$ に何を置いたかが重要となる。

ただ、具体的な値と言うよりも、「残っている中で何番目」かの情報が重要となる。
（具体的な値だと、$i-1$ 以前に何を使ったかも覚えてないといけないが、それは場合分けが多すぎて扱いきれない）

• 参考: Educational DP Contest T

以下のDPを考える。

• $DP(t)[i,j,k]=A_t,B_t$ までを決め、$A$ の末尾には残っている中で $i$ 番目に小さい値、$B$ の末尾には $j$ 番目に小さい値を用いて、暫定類似度が $k$ の順列の決め方の個数

$t$ に関しては、$DP(t)$ を求めるには $DP(t-1)$ の情報があれば十分なので、1個前のみ引き継ぐ形で実装上は省略する。

ある $t,k$ についての更新を考える。配るDPで考えると、

i\j  0 1 2 3 4 5
0
1
2          *
3
4
5

例えば、末尾を含め6個の値が残っていて、$i=2,j=3$（$A_t$ は自身より小さい値が2個残り、$B_t$ は3個残っている状態）からの遷移を考えると、

• 類似度が1増えるのは、
  • $A_{t+1}$ に $A_t$ より小さい値を、$B_{t+1}$ に $B_t$ より小さい値を置く
  • $A_{t+1}$ に $A_t$ より大きい値を、$B_{t+1}$ に $B_t$ より大きい値を置く

類似度が変わらないのはその逆なので、以下のようになる。（+が1増える、0が変わらない）

i\j  0 1 2 3 4 5
0    + + +   0 0
1    + + +   0 0
2          *
3    0 0 0   + +
4    0 0 0   + +
5    0 0 0   + +

ただし、このDPの特性上、今の値を示す行・列は次のDPでは削除されるため、次に引き継ぐ位置は以下のようになる。

i\j  0 1 2 3 4
0    + + + 0 0
1    + + + 0 0
2    0 0 0 + +
3    0 0 0 + +
4    0 0 0 + +

よって、この表を元に

• $DP(t+1)[:,:,k+1]$ のうち、“+“になっている箇所
• $DP(t+1)[:,:,k]$ のうち、”0”になっている箇所

に、$DP(t)[i,j,k]$ を加算してやればよい。

これは、1つ1つに加算していればTLEになってしまうが、更新対象のそれぞれが矩形になっているため、2次元imos法で高速化できる。

• いもす法 - いもす研 (imos laboratory)

ある矩形領域 i=[P,Q) j=[R,S) に一律 a を加算したいとき、、、

    R       S
  +-------+      4箇所に加減算した上で、
P |+a     |-a    後から2次元累積和を取ると、実現できる
  |       |
  |       |
  +-------+
Q  -a      +a

$t=N$ までやって、$DP[0,0,K]$ が答え。

Python3

import os
import sys

import numpy as np


def solve(n, kk, p):
    dp = np.zeros((n, n, kk + 2), np.int64)
    dp[:, :, 0] = 1

    for l in range(n - 1):
        rem = n - l
        ndp = np.zeros_like(dp)
        for i in range(rem):
            for j in range(rem):
                for k in range(min(kk, l) + 1):
                    a = dp[i, j, k]
                    ndp[0, 0, k + 1] += a
                    ndp[0, j, k + 1] -= a
                    ndp[i, 0, k + 1] -= a
                    ndp[i, j, k + 1] += a * 2
                    ndp[0, j, k] += a
                    ndp[i, 0, k] += a
                    ndp[i, j, k] -= a * 2
        ndp %= p
        for i in range(rem - 1):
            for j in range(rem - 1):
                ndp[i, j + 1, :] += ndp[i, j, :]
        for j in range(rem - 1):
            for i in range(rem - 1):
                ndp[i + 1, j, :] += ndp[i, j, :]
        dp = ndp % p

    return dp[0, 0, kk]


SIGNATURE = '(i8,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, k, p = map(int, input().split())
ans = solve(n, k, p)
print(ans)