AtCoder Regular Contest++ 216 A,B問題メモ

AtCoder Regular Contest++ 216

ぷらぷらが始まった（ARC Div.1に相当する新名称）。

A - Reversi 3

問題文

0 と 1 からなる長さ $N$ の文字列 $A,B$ が与えられます．$A$ の $i$ 文字目を $A_i$ とします．
あなたは以下の操作を $0$ 回以上好きな回数行うことができます．
- $A_{i-1}=A_{i+1}$ を満たす整数 $i\ (2\leq i\leq N-1)$ を選び，$A_i$ を反転する（1 ならば 0 に，0 ならば 1 にする）．
操作を繰り返すことで $A$ を $B$ に一致させることができるか判定し，可能ならば一致させるために必要な操作回数の最小値を求めてください．

$T$ 個のテストケースが与えられるので，それぞれについて答えを求めてください．

制約

$1\leq T\leq 2\times 10^5$
$3\leq N\leq 10^6$
$A,B$ は 0 と 1 からなる長さ $N$ の文字列
$T,N$ は整数
全てのテストケースに対する $N$ の総和は $10^6$ 以下

解法

まず、端の要素は変えられないので、$A_1 \neq B_1$ または $A_N \neq B_N$ の場合、不可能である。以下、同じとする。

（制約からも何となくそんな気がするとおり）貪欲っぽく左から合わせていく解法が成り立つ。
ただし、「一致させるためには $i$ で操作しないといけないが、今のままでは操作できないので、$i$ より右で一旦操作し、遡ってきて $i$ を操作できる状態にする」必要がある場合もある。

A  0 1 1 0 0 1 0
B  0 0 . . . . .
     ^------------- ここを反転させたいが、両隣が異なるので操作できない
       ^ ^ ^------- ここらも同じく無理
             v----- ここでやっと反転できる
   0 1 1 0 0 0 0
           v        一旦反転させたら、反転させられなかった箇所は
   0 1 1 0 1 0 0    順次左に1個ずつ反転させていける
         v
   0 1 1 1 1 0 0
       v
   0 1 0 1 1 0 0
     v
   0 0 0 1 1 0 0
        ↕          この時、最初のAから、
A  0 1 1 0 0 1 0    反転させたかった箇所 ～ はじめて反転できる箇所 が
     ~ ~ ~ ~ ~      まるっと反転した形となる。

その際、$A,B$ のままでは（解けなくはないが）どこから波及させられるか扱いづらいので、扱いやすい形に変換する。

以下のように $E,F$ を作る。

$C=(C_1,...,C_{N-1}),D=(D_1,...,D_{N-1})$ を、それぞれ $A,B$ の隣接要素のXORとする。
$E=(E_1,...,E_{N})$ を、$(0,C_1,...,C_{N-1},0)$ の隣接要素のXORとする。
$F=(F_1,...,F_{N})$ を、$(0,D_1,...,D_{N-1},0)$ の隣接要素のXORとする。

逆の操作をすれば $E,F$ は $A,B$ に一意に復元できるので、$E$ を $F$ に一致させればよい。

操作は $E$ 上では以下のようになる。

$E_i=0$ である $i$（$2 \le i \le N-1$）を選び、両隣の要素 $E_{i-1},E_{i+1}$ を反転させる

$i=1,2,...,N-2$ の順に、$E_i \neq F_i$ なら、$i+1$ に操作をしなければならない。

i     1 2 3 4 5 6 7 8 ...
A     0 1 1 0 0 1 0 0
B     0 0 . . . . . .
C  (0) 1 0 1 0 1 1 0
D  (0) 0 . . . . . .
E     1 1 1 1 1 0 1   ...
F     0 . . . . . .   ...

上例では、まず早速 $i=1$ の時、$E_1 \neq F_1$ のため $E_{i+1}$ に操作する必要があるが、$E_2=1$ なので今のままでは操作できない。

$j=6$ で操作可能になるので、左に派生させていくと、できるようになる。

      i         j
E     1 1 1 1 1 0 1
                v
      1 1 1 1 0 0 0
              v
      1 1 1 0 0 1 0
            v
      1 1 0 0 1 1 0
          v
      1 0 0 1 1 1 0
        v
      0 0 1 1 1 1 0

その際、操作回数は $j-i$ 回必要で、元のEからは、以下のように変化する。

$i+1=j$（操作したい $i+1$ が、そのまま $E_{i+1}=0$ だった）場合、
- $i-1,i+1$ が反転する

$i \lt j$ の場合、
- $j$ は $0→1$
- $i$ は $1→0$
- $i-1,j+1$ が反転する

この変化により、 $i+2,...,j$ は全て $1$ となるため、次に例えば $i+1$ を反転させたい場合、$i+2$ 以降で $0$ が出現するのは必ず $j+1$ 以降となる。

$i$ を進めていくのに伴い、$j$ のポインタも進めていけば、$O(N)$ で答えが求められる。

Python3

from sortedcontainers import SortedSet

def solve(n, a, b):
    aaa = list(map(int, a))
    bbb = list(map(int, b))
    if aaa[0] != bbb[0] or aaa[-1] != bbb[-1]:
        return -1
    ccc = [aaa[i] ^ aaa[i + 1] for i in range(n - 1)]
    ddd = [bbb[i] ^ bbb[i + 1] for i in range(n - 1)]
    ddd = [ccc[0] ^ 1] + ddd + [ccc[-1] ^ 1]
    ccc = [ccc[0] ^ 1] + ccc + [ccc[-1] ^ 1]
    aaa = [ccc[i] ^ ccc[i + 1] for i in range(n)]
    bbb = [ddd[i] ^ ddd[i + 1] for i in range(n)]
    idx = SortedSet(i for i, a in enumerate(aaa) if a == 0)
    ans = 0
    for i in range(n):
        if aaa[i] == bbb[i]:
            continue
        ji = idx.bisect(i)
        if ji == len(idx):
            return -1
        j = idx[ji]
        if j == n - 1:
            return -1
        aaa[i] ^= 1
        aaa[j + 1] ^= 1
        ans += j - i
        if i + 1 < j:
            aaa[i + 1] = 0
            aaa[j] = 1
            idx.remove(j)
        if aaa[j + 1] == 0:
            idx.add(j + 1)
        else:
            idx.remove(j + 1)
    return ans


t = int(input())
for _ in range(t):
    n = int(input())
    a = input()
    b = input()
    ans = solve(n, a, b)
    print(ans)

B - Range Mex Sum

問題文

正整数 $N$ および整数列 $A=(A_1,A_2,\ldots,A_N)$ が与えられます．$A_i$ は $-1$ 以上 $N-1$ 以下であり，$1\leq i\lt j\leq N$ に対し，$A_i\neq -1$ かつ $A_j\neq -1$ ならば $A_i\neq A_j$ です．
次のクエリを $Q$ 回解いてください．
- 整数 $l,r\ (1\leq l\leq r\leq N)$ が与えられる．
- $(0,1,\ldots,N-1)$ の順列 $P=(P_1,P_2,\ldots,P_N)$ であって，$A_i\neq -1\implies P_i=A_i$ を満たすようなもの全てに対する $\mathrm{mex}(\lbrace P_l,P_{l+1}, \ldots,P_r\rbrace)$ の総和を $998244353$ で割った余りを求めよ．

制約

$1\leq N\leq 5000$
$1\leq Q\leq 5\times 10^5$
$-1\leq A_i\leq N-1$
$1\leq i\lt j\leq N$ に対し，$A_i\neq -1$ かつ $A_j\neq -1$ ならば $A_i\neq A_j$
各クエリにおいて，$1\leq l\leq r\leq N$
入力は全て整数

解法

$A$ にある $-1$ の個数を $m$ とする。また、$A$ に存在しない $0～N-1$ の値を小さい順に $B_1,B_2,...,B_m$ とする。

全ての場合の総和を直接考えるより、期待値（MEXの値×そのMEXになる確率）を考えた方がわかりやすい。期待値に全体の数 $m!$ をかければ全ての場合の総和となる。

ある1クエリに答えることを考える。以下の2要素によって決まる。

$x$: その区間内にある $-1$ の個数
$k$: その区間内に無い、$-1$ でない最小の $A_i$ の値（存在しない場合は $k=N$）

例
A  _ _ 5 _ 1 3 _ _ _ 7    m=6  B=(0,2,4,6,8,9)
        [l-------r]       x=3  k=5

$B_1=0$ になるのは、$B_1$ が区間外に配置されたとき
- 確率 $\frac{m-x}{m}$
$B_2=2$ になるのは、$B_1$ が区間内、$B_2$ が区間外に配置されたとき
- 確率 $\frac{x}{m} \times \frac{m-x}{m-1}$
$B_3=4$ になるのは、$B_1,B_2$ が区間内、$B_3$ が区間外に配置されたとき
- 確率 $\frac{x}{m} \times \frac{x-1}{m-1} \times \frac{m-x}{m-2}$
$B_4 = 6 \gt k = 5$ なので、これ以上は必ず $k$ の方がMEXになる。
- $B_1,B_2,B_3$ が区間内に配置されたとき。
- 確率 $\frac{x}{m} \times \frac{x-1}{m-1} \times \frac{x-2}{m-2}$

別の例
A  _ _ 5 _ 1 3 _ 4 _ 7    m=5  B=(0,2,6,8,9)
      [l---------r]       x=2  k=7

$B_1=0$ になるのは、$B_1$ が区間外に配置されたとき
- 確率 $\frac{m-x}{m}$
$B_2=2$ になるのは、$B_1$ が区間内、$B_2$ が区間外に配置されたとき
- 確率 $\frac{x}{m} \times \frac{m-x}{m-1}$
これで区間内の $-1$ が全て埋まるので、これ以外は必ず $\min(k,B_3)$ になる。
- 確率 $\frac{x}{m} \times \frac{x-1}{m-1}$

$B_i$（$i=1,2,3,...$）がMEXになる確率は、小さい方から順に求めていける。どこかで $k$ によって打ち切られるか、$x$ が全て埋まって次の $B_{x+1}$ が確定でMEXになる。

「$k$ による打ち切り」の発生有無とその値はクエリによって異なるが、それ以外は $x$ にのみ依存する。
言い換えると、$x$ が同じクエリは、打ち切りの発生までは前計算結果を共有できる。

$k$ を無視して、以下を計算しておく。

$C'[x][i]:=$ 区間内 $-1$ が $x$ 個の時、MEXが $B_i$ になる確率
$D'[x][i]:=$ 区間内 $-1$ が $x$ 個の時、MEXが $B_i$ になる期待値（確率 $\times B_i$）
$C:=C'$ を $i$ について累積和を取ったもの
$D:=D'$ を $i$ について累積和を取ったもの

各 $x$ に対し、有効な $i$ は $1～x+1$ なので、その範囲を前計算しておく。$O(N^2)$ で計算できる。

すると、$k$ があるときの答えは、

$B$ の中で $k$ を超えない最大を $B_i$ とする。
$i=0$（$k \lt B_1$）の時、MEXは常に $k$
$x \ge i$ の時、区間内に $i$ 個置いた時点で $k$ が必ず次のMEXになる。
$x \lt i$ の時、区間内に $x$ 個置いた時点で $B_{x+1}$ が必ず次のMEXになる。

「区間内に $y$ 個置いた時点で $z$ が必ず次のMEXになる」というクエリでは、以下が期待値となる。

$D[x][y] + (1-C[x][y]) \times z$

$x$ や $k$ を求める際にセグメント木等を使うと1クエリ $O(\log{N})$、累積min・累積和などを使えば $O(1)$ で求められる。 $B,k$ から $i$ を求めるのは、二分探索で $O(\log{N})$、各 $k$ の値に対する $i$ を前計算しておけば $O(1)$ で答えられる。

全体 $O(N^2 + Q)$ または $O(N^2 + Q\log{N})$ となる。

Python3

from bisect import bisect
from operator import add

from atcoder.segtree import SegTree


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


def solve(n, q, aaa, queries):
    not_used = [True] * n
    acc_min_left = [n] * (n + 1)
    acc_min_right = [n] * (n + 1)
    sgt2 = SegTree(add, 0, n)
    for i, a in enumerate(aaa):
        if a == -1:
            sgt2.set(i, 1)
            continue
        not_used[a] = False
        acc_min_left[i + 1] = a
        acc_min_right[i] = a
    not_used = [a for a, u in enumerate(not_used) if u]
    m = len(not_used)
    not_used.append(n)

    for i in range(1, n + 1):
        acc_min_left[i] = min(acc_min_left[i - 1], acc_min_left[i])
    for i in range(n - 1, -1, -1):
        acc_min_right[i] = min(acc_min_right[i], acc_min_right[i + 1])

    MOD = 998244353
    facts, finvs = precompute_factorials(m, MOD)
    invs = [1]
    for i in range(1, m + 1):
        invs.append(finvs[i] * facts[i - 1] % MOD)

    precalc1 = [[0]]  # 区間中の-1の個数 x 毎に、(mex が not_used[i] になる確率) * (not_used[i]) の累積和
    precalc2 = [[0]]  # 区間中の-1の個数 x 毎に、(mex が not_used[i] になる確率) の累積和
    for x in range(1, m + 1):
        pc1 = [0]
        pc2 = [0]
        p = 1
        for i in range(x + 1):
            tmp = p * (m - x) * invs[m - i] % MOD
            pc1.append((pc1[-1] + tmp * not_used[i]) % MOD)
            pc2.append((pc2[-1] + tmp) % MOD)
            p = p * (x - i) * invs[m - i] % MOD
        precalc1.append(pc1)
        precalc2.append(pc2)
        # print(f'{x=} {precalc1=}')
        # print(f'{x=} {precalc2=}')
    # print(f'{not_used=}')

    buf = []
    for l, r in queries:
        l -= 1
        min1 = acc_min_left[l]
        min2 = acc_min_right[r]
        k = min(min1, min2)
        # print(f'{l=} {r=} {k=} {buf=}')

        if m == 0:
            buf.append(k)
            continue

        bi = bisect(not_used, k)
        if bi == 0:
            ans = k * facts[m] % MOD
            buf.append(ans)
            continue

        x = sgt2.prod(l, r)
        if x == 0:
            ans = not_used[0] * facts[m] % MOD
            buf.append(ans)
            continue

        if x == m:
            ans = k * facts[m] % MOD
            buf.append(ans)
            continue

        if x < bi:
            ans = precalc1[x][x]
            ans += (1 - precalc2[x][x]) * not_used[x]
            ans *= facts[m]
            ans %= MOD
            buf.append(ans)
            continue

        # print(f'{not_used=} {x=} {k=} {bi=}')
        ans = precalc1[x][bi]
        ans += (1 - precalc2[x][bi]) * k
        ans *= facts[m]
        ans %= MOD
        buf.append(ans)

    return buf


n, q = list(map(int, input().split()))
aaa = list(map(int, input().split()))
queries = [list(map(int, input().split())) for _ in range(q)]
ans = solve(n, q, aaa, queries)
print('\n'.join(map(str, ans)))

目次

AtCoder Regular Contest++ 216 A,B問題メモ

A - Reversi 3

問題文

制約

解法

B - Range Mex Sum

問題文

制約

解法