たとえば $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 \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)
単調増加列の問題は、差分に着目すると上手くいくことがよくある。
末尾と $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つを固定すると、
そのパターン数が $O(1)$ で求められる。
この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)