AtCoder Beginner Contest 276 F,G 問題メモ

AtCoder Beginner Contest 276

F - Double Chance

問題

$N$ 枚のカードがあり、カード $i$ には整数 $A_i$ が書かれている
$K=1,2,...,N$ について、次の問題を解く
問題
- カード $1～K$ の中から、等確率でカードを1枚引き、戻すことを2回繰り返す
- 1回目と2回目に引いた数をそれぞれ $x,y$ とすると、$\max(x,y)$ をその試行の得点とする
- 得点の期待値を $\mod{998244353}$ で求めよ
$1 \le N \le 2 \times 10^5$
$1 \le A_i \le 2 \times 10^5$

解法

たとえば $K=4$ で $1,4,7,10$ が書かれたカードから引く場合を考えると、全試行パターンは1回目と2回目で4通りずつなので、$K^2=16$ 通り。

期待値は、全ての場合の得点を合計してこの $K^2$ で割ればよいので、合計の方を求められればよい。

得点ごとに、そのパターン数を考えると

得点が1になる場合:  1を2回連続で引く              1x1       = 1 通り
得点が4になる場合:  1,4から2回連続で引く場合から
                    得点が1になる場合を除く       2x2 - 1x1 = 3 通り
得点が7になる場合:  3枚から2回連続で引く場合から
                    得点が1,4になる場合を除く     3x3 - 2x2 = 5 通り
得点が10になる場合: 4枚から2回連続で引く場合から
                    得点が1,4,7になる場合を除く   4x4 - 3x3 = 7 通り

1x1 + 4x3 + 7x5 + 10x7 = 118

小さい順に $1,3,5,...$（奇数）通りずつパターン数が増えていく。得点とパターン数を掛け合わせればよい。
（同じ数字があった場合も、適当に順番を決めて扱えばよい）

で、この $K=4$ の値から $K=5$ の値を求められる。

$A_5=6$ を追加した場合、

  得点 パターン数
     1      1
     4      3
New! 6      5
     7   5→7
    10   7→9

3番目に追加されるのでパターン数は5となり、$6 \times 5$ が新たに増える。
また、それより大きい値の $7,10$ は、1つにつき $2$ ずつパターン数が増える。

なので、

追加する $A_i$ より小さい値が何個あるか $p$
追加する $A_i$ 以上となる値の合計はいくらか $q$

が、逐次求められればよいことになる。（$A_i \times (2p+1) + 2q$ だけ増える）

これは、$A_i$ をindexとして個数を管理するFenwick Treeと、合計を管理するFenwick Treeの2本を使って 1回あたり $\log{A_{max}}$ で求められる。

Python3

from operator import add


class FenwickTreeInjectable:
    def __init__(self, n, identity_factory, func):
        self.size = n
        self.tree = [identity_factory() for _ in range(n + 1)]
        self.func = func
        self.idf = identity_factory
        self.depth = n.bit_length()

    def add(self, i, x):
        i += 1
        tree = self.tree
        func = self.func
        while i <= self.size:
            tree[i] = func(tree[i], x)
            i += i & -i

    def sum(self, i):
        i += 1
        s = self.idf()
        tree = self.tree
        func = self.func
        while i > 0:
            s = func(s, tree[i])
            i -= i & -i
        return s

    def debug_print(self):
        prev = 0
        arr = []
        for i in range(self.size):
            curr = self.sum(i)
            arr.append(curr - prev)
            prev = curr
        print(arr)


n = int(input())
aaa = list(map(int, input().split()))
MOD = 998244353

limit = 200_001
fwt_count = FenwickTreeInjectable(limit, int, add)
fwt_value = FenwickTreeInjectable(limit, int, add)
total = 0
ans = 0
for i, a in enumerate(aaa, start=1):
    c = fwt_count.sum(a)
    v = fwt_value.sum(a)
    fwt_count.add(a, 1)
    fwt_value.add(a, a)
    ans += (2 * c + 1) * a
    ans += (total - v) * 2
    total += a
    print(ans * pow(i * i, MOD - 2, MOD) % MOD)

G - Count Sequences

問題

以下の条件を満たす項数 $N$ の整数列 $A=(a_1,a_2,...,a_N)$ の個数を $\mod{998244353}$ で求めよ
条件
- 要素は $0$ 以上 $M$ 以下で広義単調増加 $0 \le a_1 \le a_2 \le ... \le a_N \le M$
- 隣り合う項（$i=1～N-1$ に対し、$a_i$ と $a_{i+1}$）は、3で割ったあまりが異なる
$2 \le N \le 10^7$
$1 \le M \le 10^7$

解法

単調増加列の問題は、差分に着目すると上手くいくことがよくある。
末尾と $M$ との差分（下図 $b_5$）まで含めて考えると、差分の総和が上限値 $M$ となる。

A        0   a1   a2   a3   a4   M
Aの差分   b1 + b2 + b3 + b4 + b5    = M

隣り合う項の3で割ったあまりが異なるので、両端を除く $b_2～b_4$ は、3の倍数であってはいけないということになる。

逆にこれさえ満たせば条件を満たす数列となるため、そのような非負整数列 $B=(b_1,...,b_{N+1})$ の個数を数えればよい。

ここで、$B$ を3で割った商 $X=(x_1,...,x_{N+1})$ とあまり $Y=(y_1,...,y_{N+1})$ に分けて考える。

条件は「$y_2～y_N$ は $0$ 禁止（$1,2$ のいずれか）」「$3x_i+y_i$ の総和が $M$」となる。

条件を満たす $X,Y$ は、例えば以下のようにして構築できる。

$y_1$ を $0,1,2$ の中から決める
$y_2～y_N$ は最低でも1なのであらかじめ1を敷いておき、「2になる要素がどれか」を決める
ここまでに決めた数によって、$M$ の残りから、$X$ の総和と $y_{N+1}$ は自動的に決まる
$X$ の総和の $x_1,...,x_{N+1}$ への割り振り方を決める

以下の2つを固定すると、

$s$: $y_1$ の値（$0/1/2$ の3通り）
$k$: $y_2～y_N$ のうち、$2$ になる要素の個数（$0～N-1$ の $N$ 通り）

そのパターン数が $O(1)$ で求められる。

$y_2～y_N$ のうちの $2$ の配置の決め方は ${}_{N-1}C_{k}$
$X$ の総和は $\dfrac{M-(N-1)-s-k}{3}$（切り捨て）に決まる。これを $t$ とする
- $x_1,...,x_{N+1}$ への割り振り方は ${}_{N+1}H_{t}$

この2つを掛け合わせたものが、$s,k$ を固定した時のパターン数となる。全探索して合計すれば答え。

本番では3で割った商とあまりに分割するところまでは思いついたが、$b_1$ を特別視しすぎて $b_2～b_N$ だけで $X,Y$ を考え、$b_1$ は別で計算してしまったため、考慮すべきパターン数が増え、計算量が $NM$ から減らせなかった。

Python3

import os
import sys

import numpy as np


def solve(n, m):
    def mod_pow(x, a, MOD):
        ret = 1
        cur = x
        while a:
            if a & 1:
                ret = ret * cur % MOD
            cur = cur * cur % MOD
            a >>= 1
        return ret

    def prepare_factorials(n, MOD):
        factorials = np.ones(n + 1, np.int64)
        for m in range(1, n + 1):
            factorials[m] = factorials[m - 1] * m % MOD
        inversions = np.ones(n + 1, np.int64)
        inversions[n] = mod_pow(factorials[n], MOD - 2, MOD)
        for m in range(n, 1, -1):
            inversions[m - 1] = inversions[m] * m % MOD
        return factorials, inversions

    MOD = 998244353

    facts, finvs = prepare_factorials(n + m, MOD)

    def ncr(n, r):
        if r < 0 or n < r: return 0
        return facts[n] * finvs[r] % MOD * finvs[n - r] % MOD

    def nhr(n, r):
        if r < 0: return 0
        return facts[n + r - 1] * finvs[r] % MOD * finvs[n - 1] % MOD

    ans = 0

    for a in range(n + 2):
        rem = m - (n - 1) - a
        if rem < 0:
            break

        base = (ncr(n - 1, a) + ncr(n - 1, a - 1) + ncr(n - 1, a - 2)) % MOD
        base *= nhr(n + 1, rem // 3)

        ans += base
        ans %= MOD

    return ans


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

目次

AtCoder Beginner Contest 276 F,G 問題メモ

F - Double Chance

問題

解法

G - Count Sequences

問題

解法