AtCoder Beginner Contest 246 Ex問題メモ

AtCoder Beginner Contest 246

Ex - 01? Queries

問題

'0','1','?' の3種の文字のみからなる長さ $N$ の文字列 $S$ が与えられる
$Q$ 個のクエリを順番に処理せよ
クエリ
- $x_i,c_i$ が与えられる
- $S$ の $x_i$ 文字目を $c_i$ （'0','1','?' のいずれか）に書き換える
  - この変更は以降のクエリにも反映される
- 「 $S$ に含まれる全ての '?' を '0' または '1' に独立に置き換えて得られる文字列の、（連続とは限らない）部分列」としてあり得る文字列の個数を $\mod{998244353}$ で出力する
$1 \le N,Q \le 10^5$
実行時間制限: 6sec.

解法

DPの遷移が、小さな行列の積で表現できることを利用して、行列を乗せたセグメント木で必要な箇所だけを更新する。

公式Editorialでは3×3の行列を用いていた。その方が自然なので、基本的にそちらを参照すればよい。

4×4の行列を使って解いたので、考えだけメモ。

$DP[i][j]=S$ の先頭 $i$ 文字で作れる部分列で、( $j:$ 後述の分類)に属するものの個数

$j$ は、以下の4つの状態の分類とする。

oo: 後ろに'0'を置いても'1'を置いても、先頭 $i$ 文字からは作れない部分列となる
ox: 後ろに'0'を置いたら新しい部分列となるが、'1'を置いた部分列は先頭 $i$ 文字からも作れる
xo: 後ろに'1'を置いたら新しい部分列となるが、'0'を置いた部分列は先頭 $i$ 文字からも作れる
xx: 後ろに'0'を置いても'1'を置いても、先頭 $i$ 文字から既に作れる部分列となる

それぞれの状態から、 $S$ の $i+1$ 文字目が '0','1','?' のそれぞれの場合だったときに、「それを末尾に採用したとき（採用したら新しい文字列となる場合のみ）」「採用しなかったとき」で、遷移先が一意に定まる。

   0                1                ?                    遷移先
   oo ox xo xx      oo ox xo xx      oo ox xo xx          ------→
oo [1, 0, 1, 0],    [1, 1, 0, 0],    [2, 0, 0, 1],    遷｜
ox [1, 0, 0, 1],    [0, 1, 0, 0],    [1, 0, 0, 1],    移｜
xo [0, 0, 1, 0],    [1, 0, 0, 1],    [1, 0, 0, 1],    元｜
xx [0, 0, 0, 1],    [0, 0, 0, 1],    [0, 0, 0, 1],      ↓

大まかには、「先頭 $i$ 文字から作れる部分列の集合」を $B_i$ として、

$S_{i+1}$ を末尾に採用して新規部分列としたら、その後ろに何かをさらにくっつけた部分列は $B_{i+1}$ には含まれないので、ooに遷移できる（遷移元の状態問わず）
- たとえば $B_i$ に '00' が含まれ、 $S_{i+1}=1$ を追加して新規部分列 '001' が作れるとなったとき、 $B_{i+1}$ に '0010' や '0011' が既に含まれる、という状態なら、 $B_i$ の時点で '001' が既に含まれていないとおかしいので矛盾
$S_{i+1}$ を採用できるのに採用しなかったら、採用したものも $B_{i+1}$ に含まれるので、次からは採用できなくなる
- $S_{i+1}=0$ のとき、oo なら xo に遷移、ox なら xx に遷移、など

あとは $S$ の各文字に対応する行列を先頭からかけ算していった結果で、1行目の和が答えとなる（初期状態の空文字列は oo にあたるので）。
ただし結果には空文字列も含まれるため、最後に -1 する。

セグメント木の実装に注意。行列積は左項と右項の順番で答えが変わるので、ここを蔑ろにした実装では正しい答えが求まらない。

行列積の計算量はサイズの3乗で利いてくるので、 $3^3=27$ と $4^3=64$ では2倍以上効率が悪いが、Pythonでもギリギリで通った。

Python3

        
          
          
              
              import numpy as np
 
 
class SegmentTreeInjectable:
    def __init__(self, n, identity_factory, func):
        n2 = 1 << (n - 1).bit_length()
        self.offset = n2
        self.tree = [identity_factory() for _ in range(n2 << 1)]
        self.func = func
        self.idf = identity_factory
 
    @classmethod
    def from_array(cls, arr, identity_factory, func):
        """ 既存の配列から生成 """
        ins = cls(len(arr), identity_factory, func)
        ins.tree[ins.offset:ins.offset + len(arr)] = arr
        for i in range(ins.offset - 1, 0, -1):
            l = i << 1
            r = l + 1
            ins.tree[i] = func(ins.tree[l], ins.tree[r])
        return ins
 
    def add(self, i, x):
        """
        Aiにxを加算
        :param i: index (0-indexed)
        :param x: add value
        """
        i += self.offset
        self.tree[i] = self.func(self.tree[i], x)
        self.__upstream(i)
 
    def update(self, i, x):
        """
        Aiの値をxに更新
        :param i: index(0-indexed)
        :param x: update value
        """
        i += self.offset
        self.tree[i] = x
        self.__upstream(i)
 
    def __upstream(self, i):
        tree = self.tree
        func = self.func
        while i > 1:
            j = i >> 1
            tree[j] = func(tree[j << 1], tree[(j << 1) | 1])
            i >>= 1
 
    def get_range(self, a, b):
        """
        [a, b)の値を得る
        :param a: index(0-indexed)
        :param b: index(0-indexed)
        """
        tree = self.tree
        func = self.func
        result_l = self.idf()
        result_r = self.idf()
 
        l = a + self.offset
        r = b + self.offset
        while l < r:
            if r & 1:
                result_r = func(tree[r - 1], result_r)
            if l & 1:
                result_l = func(result_l, tree[l])
                l += 1
            l >>= 1
            r >>= 1
 
        return func(result_l, result_r)
 
    def get_all(self):
        return self.tree[1]
 
    def get_point(self, i):
        return self.tree[i + self.offset]
 
    def debug_print(self):
        i = 1
        while i <= self.offset:
            print(self.tree[i:i * 2])
            i <<= 1
 
 
mat_zero = np.array([
    [1, 0, 1, 0],
    [1, 0, 0, 1],
    [0, 0, 1, 0],
    [0, 0, 0, 1],
], np.int64)
mat_one = np.array([
    [1, 1, 0, 0],
    [0, 1, 0, 0],
    [1, 0, 0, 1],
    [0, 0, 0, 1],
], np.int64)
mat_ques = np.array([
    [2, 0, 0, 1],
    [1, 0, 0, 1],
    [1, 0, 0, 1],
    [0, 0, 0, 1],
], np.int64)
mat_identity = np.eye(4, dtype=np.int64)
 
MOD = 998244353
 
n, q = map(int, input().split())
s = input()
init_array = []
for c in s:
    if c == '0':
        init_array.append(mat_zero)
    elif c == '1':
        init_array.append(mat_one)
    else:
        init_array.append(mat_ques)
 
sgt = SegmentTreeInjectable.from_array(init_array, lambda: mat_identity, lambda a, b: a @ b % MOD)
 
for _ in range(q):
    x, c = input().split()
    x = int(x) - 1
    if c == '0':
        sgt.update(x, mat_zero)
    elif c == '1':
        sgt.update(x, mat_one)
    else:
        sgt.update(x, mat_ques)
    res = sgt.get_all()
 
    ans = (res[0].sum() - 1) % MOD
    print(ans)

            

        
      