東京海上日動プログラミングコンテスト2023（AtCoder Beginner Contest 299）F問題メモ

東京海上日動プログラミングコンテスト2023（AtCoder Beginner Contest 299）

F - Square Subsequence

• 英小文字からなる文字列 $S$ が与えられる
• 以下の条件を満たす文字列 $T$ の個数を $998244353$ で割ったあまりで求めよ
  • $T$ を2つ連結した文字列 $TT$ が、$S$ の（連続とは限らない）部分列として含まれる
• $1 \le |S| \le 100$

$S$ を適当なところで分割して（$S_1,S_2$）、それぞれから同じ文字列が部分列として取れたら、それは $T$ の条件を満たす。

ababbaba
  ↓
ababb | aba  適当な場所で分割
  ↓
a  b  | ab   などが共通して取れる → "ab" はTの1つである

2つの文字列の共通部分列の個数は、2次元DPで数えられる。
たとえば「$S_1,S_2$ のどちらかでも違う箇所から取り出した文字列は、文字列として同じであっても違うものとして数える」場合は、 以下のように行うことができる。

• $DP[i][j]=S_1[0...i]$ と $S_2[0...j]$ を使って、特に $S_1[i]$ と $S_2[j]$ は必ず使ってできる、共通部分列の個数

      S2
      -  a  b  a
S1 -  1              S1のi文字目とS2のj文字目が等しい場合に限り、
   a     1     1     DPの 0～i-1 行目、0～j-1 列目の矩形の総計が
   b        2        DP[i][j] になる
   a     1     4
   b        3
   b        3

一方、この問題は、同じ文字列 $T$ は1回しか数えてはいけない。

なので、重複を除くときの定石として「同じ文字列なら、必ず最も左詰で取るものだけ数える」という条件が必要になる。

• $DP[i][j]=S_1[0...i]$ と $S_2[0...j]$ を使って、特に $S_1[i]$ と $S_2[j]$ は必ず使ってできる、それぞれから左詰で取った共通部分列の個数

どうすれば、最も左詰のものだけ数えたことになるか？
より具体的には、

■S2[0]は必ず使う

  ababb | aba から ba が見つかったとして、区切り位置を変えることで、
   ~~      ~~
  ababba | ba など全く同じ位置から ba が拾えてしまう。
   ~~      ~~
  S2の方の共通部分列は、必ず先頭から始まるとすることで、
  そこが区切りとなった1回のみで数えられるようにできる。
  
  
  ここからDPの初期状態を考えると、
  S2[0] = c0 として、S1 の中で最も左にある c0 が対応する。
  その位置を i0 として、DP[i0][0] = 1 が初期状態となる。


■DP[i][j]から末尾に文字 c をつける際、最も近い c を使う
  
  DPの更新を考える。
  もらうDPより配るDPの方が重複を除くのは楽なので、それで考える。
  既に見つかっている共通部分列から、
  それぞれの末尾に文字 c を付けることを考える。（この c を26通り試す）
  
           *             *
  S1:abcdefggg と S2:cdefggg で、"e" を最後に使った共通部分列から、
  次に "g" を付ける、となったとき *印をつけたgのみに遷移し、その後のgには遷移しないようにする。
  
  f(i,c) = 位置 i から見て、次に出てくる文字 c のindex
  を事前計算しておくとやりやすい。

これで、DPは重複を除いたものが作れる。ただしこの中から答えに加算できるのは一部のみで、

■S2[0] = c0 が S1 の中で最後に出てきた位置を k として、
  答えに加算されるのは DP[k 以降][:] の総和
  
  abcabca | abcabc から左詰で "abcabc" が見つかったとして、実は
  ~~~~~~    ~~~~~~
  abcabc | aabcabc と分割したときに、既に同じものが見つかっている。
  ~~~~~~   ~ ~~~~~
  
  「左詰」にしなければならないのはS1とS2の間にも適用されて、
  S1から"abcabc"と左詰で取った、次の c0 が、S2[0]でなければならない。
  
  よって、S1の中で最後の c0 以降の文字が、共通部分列に必ず使われていないといけない。

これらを考慮してDPを実装することで、重複を除いた値が数えられる。

Python3

import os
import sys

import numpy as np


def solve(sss):
    n = len(sss)
    MOD = 998244353

    next_pos = np.full((26, n + 2), n, np.int64)
    for i in range(n - 1, -1, -1):
        c = sss[i]
        next_pos[:, i] = next_pos[:, i + 1]
        next_pos[c, i] = i

    prev_pos = np.full((26, n + 2), -1, np.int64)
    for i in range(n):
        c = sss[i]
        prev_pos[:, i + 2] = prev_pos[:, i + 1]
        prev_pos[c, i + 2] = i

    ans = 0

    for j0 in range(1, n):
        m1 = j0
        m2 = n - j0

        i0 = next_pos[sss[j0], 0]
        if i0 >= j0:
            continue

        dp = np.zeros((m1, m2), dtype=np.int64)
        dp[i0, 0] = 1

        for i in range(i0, m1):
            for _j in range(0, m2):
                j = _j + j0
                if dp[i, _j] == 0:
                    continue
                for c in range(26):
                    ni = next_pos[c, i + 1]
                    nj = next_pos[c, j + 1]
                    if ni >= j0 or nj == n:
                        continue
                    _nj = nj - j0
                    dp[ni, _nj] += dp[i, _j]
                    dp[ni, _nj] %= MOD

        pj = prev_pos[sss[j0], j0 + 1]
        ans += dp[pj:, :].sum()
        ans %= MOD

        # print(j0, pj, dp[pj:, :].sum())
        # print(dp)

    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.fromiter(map(ord, input()), dtype=np.int64) - ord('a')
ans = solve(inp)
print(ans)