AtCoder Beginner Contest 473 D,E,F,G問題メモ

D - Coefficient Stair

問題文

非負整数からなる長さ $N$ の列 $A=(A _ 1,A _ 2,\ldots,A _ N)$ であって、$\displaystyle\sum _ {i=1} ^ Ni\times A _ i=K$ を満たすものを辞書順で小さいほうからすべて出力してください。
ここで、条件を満たす数列が $3\times10 ^ 5$ 個以下であるような入力のみが与えられます。

制約

$1\le N\le 10$
$1\le K\le 2\times10 ^ 5$
条件を満たす数列は $3\times10 ^ 5$ 個以下
入力はすべて整数

解法

個人的な好みではあるんだが、問題を解くのって「入力に対し未知の出力結果を知りたいからプログラムを作る」感覚があるので、この問題のように「その未知の結果（条件を満たす数列の個数）を知らないと設定できない入力制約」って、循環的でどことなく奇妙に感じてしまう。

$K$ の範囲は大きいが、$K$ が増えると爆発的に条件を満たす数列の個数が増えるので、実質的な上限はさほど大きくないことが想定される。例えば $N=10$ の場合、実質 $K \le 64$ という制約となる。（この時 $296320$ 個）

このことに気付けば、比較的愚直な解法が可能であると推測できる。以下の再帰関数でDFSすれば求められる。

$f(tmp, i, j):=$ 現在 $tmp=(A_1,...,A_{i-1})$ と決め、次に $A_i$ を決める段階で、残りの和を $j$ とするような方法を列挙

$A_i$ を $t=0～\left \lfloor \dfrac{j}{i} \right \rfloor$ まで試しながら $f((tmp...,t),i+1,j-it)$ に進めていくとよい。

Python3

def solve(n, k):
    ans = ['0']
    ans.clear()

    def sub(cur, i, k):
        if i == n:
            d, m = divmod(k, i)
            if m == 0:
                cur.append(d)
                ans.append(' '.join(map(str, cur)))
                cur.pop()
            return
        if 0 < k < i:
            return
        for a in range(k // i + 1):
            cur.append(a)
            sub(cur, i + 1, k - i * a)
            cur.pop()

    cur = [0]
    cur.clear()
    sub(cur, 1, k)
    return ans


n, k = list(map(int, input().split()))
ans = solve(n, k)
print('\n'.join(ans))

E - K-Divisible Subarrays

問題文

非負整数からなる長さ $N$ の列 $A=(A _ 1,A _ 2,\ldots,A _ N)$ が与えられます。
非負整数列からなる長さ $1$ 以上の列 $S=(S _ 1,S _ 2,\ldots,S _ k)$ のスコアを、列 $S _ i\ (1\le i\le k)$ のうち要素の総和が $K$ で割り切れるものの個数として定めます。
$A$ を $1$ つ以上の連続する部分列に分割して得られる非負整数列の列に対するスコアの最大値を求めてください。

制約

$1\le N\le2\times10 ^ 5$
$1\le K\le10 ^ 9$
$0\le A _ i\lt K$
入力はすべて整数

解法

「$B=$ 累積和 $\bmod{K}$」を計算しておく。

K=10
i     1  2  3  4  5  6  7  8  9 10 11 12 13 14 15 16 17
A     3  1  4  1  5  9  2  6  5  3  5  8  9  7  9  3  2
B  0  3  4  8  9  4  3  5  1  6  9  4  2  1  8  7  0  2

すると、$A$ から取れる部分列のうち、総和が $K$ で割り切れる中で「最も右端が左」なのは $(4,1,5)$ とわかる。
$B$ の中で“2個目の登場が最も早い値”が $i=2,5$ の $4$ なので。

この時 $(4,1,5)$ は必ず使うとして損しない。仮に最適解の1つにおいて、これより右端が右にある部分列が “最初にスコアに寄与する部分列” なら、スコアを減らさず $(4,1,5)$ を使うように変更できるので。

同様に、その次にスコアに寄与するのは $(4,1,5)$ 以降で $B$ の中で同じ値が2回登場する最も早い区間としてよい。
上例だと、$i=5$ から再スタートして $i=5,11$ の $4$ が最も早いので、$(9,2,6,5,3,5)$ となる。

更にその次は $i=12,17$ の $2$ なので $(9,7,9,3,2)$、、、と貪欲に決めていける。

Python3

def solve(n, k, aaa):
    ans = 0
    acc = 0
    appeared = {0}
    for a in aaa:
        acc = (acc + a) % k
        if acc in appeared:
            ans += 1
            appeared.clear()
        appeared.add(acc)
    return ans


n, k = list(map(int, input().split()))
aaa = list(map(int, input().split()))
ans = solve(n, k, aaa)
print(ans)

F - A/AB Insertion

問題文

文字列 $T$ が「良い文字列」であるとは、$T$ が以下の操作で得られる文字列であることを指します。
- 空文字列から始め、以下の $2$ つの操作を好きな順序で $0$ 回以上何回でも行う。
  - 文字列の(先頭・末尾を含めた)任意の箇所を選択し、そこに A を挿入する。
  - 文字列の(先頭・末尾を含めた)任意の箇所を選択し、そこに AB を挿入する。
長さ $N$ の A, B からなる文字列 $S$ が与えられます。
以下のクエリを合計 $Q$ 個処理してください。
- 1 i c :
  - $S$ の $i$ 文字目を $c$ に変更する。
- 2 l r :
  - 現在の文字列 $S$ の $l$ 文字目から $r$ 文字目までを抜き出した文字列が「良い文字列」であれば Yes 、そうでなければ No と出力する。

制約

$N$ は $1 \le N \le 5 \times 10^5$ を満たす整数
$S$ は A, B からなる長さ $N$ の文字列
$Q$ は $1 \le Q \le 2 \times 10^5$ を満たす整数
与えられるクエリはタイプ $1,2$ のいずれかである
タイプ $1$ のクエリは以下の制約を満たす
- $i$ は $1 \le i \le N$ を満たす整数
- $c$ は A, B のいずれか
タイプ $2$ のクエリは以下の制約を満たす
- $l,r$ は $1 \le l \le r \le N$ を満たす整数

解法

E問題ほどあからさまではないが、これも言い換えると同じく累積和が関わってくる。

良い文字列であるとは、以下と同義である。

'A'を $+1$、'B'を $-1$ に置換して、先頭から累積和を取った時に、一度も負値にならない

'B'は挿入するとき、その前に必ず'A'を伴わないといけない。つまり'B'の前にはその個数以上の'A'がないといけない。
逆にそれが満たされるなら、操作を逆に考えて「'A' または 'AB' を取り出していく」と考えたとき、 'B'がなくなるまで'AB'の並びは常に存在するので、構築可能である。

よって、区間加算・区間最小値取得の遅延セグメント木などで可能である。

累積和を載せておき、クエリ1でA→Bに変わったならそれ以降に全て $-2$、 B→Aなら $+2$ を加算すれば累積和が正しい状態に保たれる。

クエリ2では累積和の「$l$ の値」と「$[l,r]$ の最小値」が等しいか調べればよい。

Python3

from itertools import accumulate
from operator import add

from atcoder.lazysegtree import LazySegTree


def solve(n, s, q, queries):
    aaa = [1 if c == 'A' else -1 for c in s]
    init = [0] + list(accumulate(aaa))
    lst = LazySegTree(min, 1 << 60, add, add, 0, init)

    buf = []
    for query in queries:
        if query[0] == '1':
            i = int(query[1]) - 1
            c = 1 if query[2] == 'A' else -1
            if aaa[i] == c:
                continue
            aaa[i] = c
            lst.apply(i + 1, n + 1, c * 2)
        else:
            l = int(query[1])
            r = int(query[2])
            x = lst.get(l - 1)
            y = lst.prod(l - 1, r + 1)
            if x == y:
                buf.append('Yes')
            else:
                buf.append('No')
    return buf


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

G - Wipeout

問題文

表に整数 $1,2,\dots,N$ が書かれたカードが $1$ 枚ずつ、計 $N$ 枚のカードが裏返しで一列に並べられています。
- これらのカードの並び順は $N!$ 通りの中から一様ランダムに定められます。
あなたは表に整数 $1,2,\dots,N$ が書かれたカードが $1$ 枚ずつあること、カードが一様ランダムに並べられたことを知っていますが、カードの表に書かれている整数に関する情報はこれ以上持っていません。
あなたは以下のゲームを行います。
- 最初、変数 $x=1$ とする。
- $x \le N$ である限り、以下の操作を繰り返す。 $1$ 回の操作は以下の $3$ つの手順からなる。
  - カードを $1$ 枚指定して、それを表に向ける。
  - カードに書かれた整数が $x$ ならそのカードを食べ、 $x$ に $1$ 加算する。
  - そうでないなら、そのカードを裏に向ける。あなたはそのカードに書かれた整数を永久に記憶しておくことができる。
あなたは常に全てのカードを食べ終わるまでの操作の総数の期待値が最小となるように行動します。
この場合、操作の総数が $K$ 回となる確率はいくらでしょうか？ $\bmod{998244353}$ で求めてください。

制約

入力は全て整数
$1 \le N \le 5 \times 10^5$
$N \le K \le 10^9$

解法

先頭から順にめくるものと考えてよい。
カードの並び順を $P$、めくる順を $Q$ のように2種類の $1～N$ の順列 $P,Q$ を定めたとして、それぞれは一様ランダムであり、結局その2つを合成した順列を先頭からめくるのと同じである。

まず、$1$ が出ないと進まない。
それまでに出たカードは、$x$ がその値になったときにそれぞれもう一度だけめくることになる。それまではめくる必要は無い。

$1$ が出たら、それまでに $2,3,...$ が出ていたらそれをめくり、再び $1$ の次からめくっていないカードをめくっていく。次に出てくれないと進まない値は「まだ出てない最小値」となる。

これを繰り返し最後までめくると、経験した「出てくれないと進まない値（★）」の個数を $L$ として、 $L$ 枚を1度だけめくり、$N-L$ 枚は都合2回ずつめくることになるので、操作回数は $2N-L$ となる。

$L$ の値は $L=2N-K$ より入力から決まる。これが $1 \le L \le N$ の範囲にならない場合は確率 $0$。

で、$L$ に対する答えを求めていく。
最後に $N!$ で割るとして、★の個数が $L$ 個となるような順列の個数が分かればいい。これは順列を $1,2,...,N$ の順に挿入して構築していく中で、「末尾に追加した回数が $L$ 回」であるようなものの個数と言い換えられる。

よって、$O(N^2)$ でTLEとはなるが、以下のようなDPで求めることができる。

$\mathrm{DP}[i,j]:=i$ までを挿入し、末尾に挿入した回数が $j$ 回であるような、$1～i$ の順列の個数

$\mathrm{DP}[0,0]=1$ より始め、$\mathrm{DP}[N,L]$ が求めたいものとなる。

このDPは、$i$ を「末尾に挿入するなら1通り」「末尾以外に挿入するなら $i-1$ 通り」の挿入箇所があるので、

i\j     0   1   2   3   4
0       1
     x0↓↘x1
1       0   1
     x1↓↘↓↘x1
2       0   1   1
     x2↓↘↓↘↓↘x1
3       0   2   3   1
     x3↓↘↓↘↓↘↓↘x1
4       0   6  11   6   1

このように、↓は $i-1$ 倍、↘は $1$ 倍で $\mathrm{DP}[i-1]$ の結果を遷移させればよい。

これは形式的冪級数に置き換えると $(i-1 + x)$ をかけることに相当する。（A130534）

よって、全体では $x(x+1)(x+2)...(x+N-1)$ の結果の $x^L$ の係数を求めればよいということになる。

項数の小さい方から畳み込んでいくことで、$O(N (\log{N})^2)$ に高速化できる。

Python3

class NttConvolution:
    """
    NTT (Number Theoretic Transform) for Codon.
    MOD = 998244353 = 119 * 2^23 + 1, primitive root = 3.
    Maximum transform length: 2^23 = 8388608.

    Usage:
        ntt = NttConvolution()
        c = ntt.convolution(a, b)
        c = ntt.convolution_all([a, b, c, d])
        ntt.ntt(a)   # in-place forward transform
        ntt.intt(a)  # in-place inverse transform
    """

    _roots: list[int]
    _inv_roots: list[int]

    def __init__(self) -> None:
        _M = 998244353
        _G = 3
        _LOG2 = 23
        roots = [0] * (_LOG2 + 1)
        inv_roots = [0] * (_LOG2 + 1)
        roots[_LOG2] = NttConvolution._pow_mod(_G, (_M - 1) // (1 << _LOG2))
        inv_roots[_LOG2] = NttConvolution._pow_mod(roots[_LOG2], _M - 2)
        for i in range(_LOG2 - 1, -1, -1):
            roots[i] = roots[i + 1] * roots[i + 1] % _M
            inv_roots[i] = inv_roots[i + 1] * inv_roots[i + 1] % _M
        self._roots = roots
        self._inv_roots = inv_roots

    @staticmethod
    def _pow_mod(base: int, exp: int) -> int:
        _M = 998244353
        result = 1
        base %= _M
        while exp > 0:
            if exp & 1:
                result = result * base % _M
            exp >>= 1
            base = base * base % _M
        return result

    @staticmethod
    def _bit_length(n: int) -> int:
        k = 0
        while n > 0:
            k += 1
            n >>= 1
        return k

    def _transform(self, a: list[int], invert: bool) -> None:
        _M = 998244353
        n = len(a)
        k = NttConvolution._bit_length(n) - 1
        table = self._inv_roots if invert else self._roots

        j = 0
        for i in range(1, n):
            bit = n >> 1
            while j & bit:
                j ^= bit
                bit >>= 1
            j ^= bit
            if i < j:
                a[i], a[j] = a[j], a[i]

        length = 1
        for s in range(k):
            length <<= 1
            half = length >> 1
            w = table[s + 1]
            i = 0
            while i < n:
                wn = 1
                for t in range(half):
                    u = a[i + t]
                    v = a[i + t + half] * wn % _M
                    a[i + t] = (u + v) % _M
                    a[i + t + half] = (u - v + _M) % _M
                    wn = wn * w % _M
                i += length

        if invert:
            inv_n = NttConvolution._pow_mod(n, _M - 2)
            for i in range(n):
                a[i] = a[i] * inv_n % _M

    def ntt(self, a: list[int]) -> None:
        self._transform(a, False)

    def intt(self, a: list[int]) -> None:
        self._transform(a, True)

    def convolution(self, a: list[int], b: list[int]) -> list[int]:
        _M = 998244353
        n = len(a)
        m = len(b)
        if n == 0 or m == 0:
            return []
        result_len = n + m - 1

        if min(n, m) <= 60:
            c = [0] * result_len
            if n < m:
                for i in range(n):
                    ai = a[i]
                    if ai == 0:
                        continue
                    for j in range(m):
                        c[i + j] = (c[i + j] + ai * b[j]) % _M
            else:
                for j in range(m):
                    bj = b[j]
                    if bj == 0:
                        continue
                    for i in range(n):
                        c[i + j] = (c[i + j] + a[i] * bj) % _M
            return c

        sz = 1
        while sz < result_len:
            sz <<= 1

        fa = [0] * sz
        fb = [0] * sz
        for i in range(n):
            fa[i] = a[i] % _M
        for i in range(m):
            fb[i] = b[i] % _M

        self._transform(fa, False)
        self._transform(fb, False)

        for i in range(sz):
            fa[i] = fa[i] * fb[i] % _M

        self._transform(fa, True)
        return fa[:result_len]

    def convolution_all(self, polys: list[list[int]]) -> list[int]:
        from collections import deque
        if len(polys) == 0:
            return [1]
        q: deque[list[int]] = deque(polys)
        while len(q) > 1:
            a = q.popleft()
            b = q.popleft()
            q.append(self.convolution(a, b))
        return q[0]


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, k):
    l = n * 2 - k
    if not (1 <= l <= n):
        return 0

    ntt = NttConvolution()
    res = ntt.convolution_all([[i, 1] for i in range(1, n)])
    # print(res)
    ans = res[l - 1]

    MOD = 998244353
    _, finvs = precompute_factorials(n, MOD)
    ans *= finvs[n]
    ans %= MOD
    if ans < 0:
        ans += MOD

    return ans


n, k = list(map(int, input().split()))
ans = solve(n, k)
print(ans)

目次

AtCoder Beginner Contest 473 D,E,F,G問題メモ

D - Coefficient Stair

問題文

制約

解法

E - K-Divisible Subarrays

問題文

制約

解法

F - A/AB Insertion

問題文

制約

解法

G - Wipeout

問題文

制約

解法