https://smijake3.hatenablog.com/entry/2018/06/13/004342

CodeChef June Challenge 2018の問題: Expected Buildings (Code: BUILDIT)
問題ページ: https://www.codechef.com/JUNE18A/problems/BUILDIT

問題概要

1つの円があり、円の中心にChefが立っている。
この円を周囲は ${h}$ 個の区間に分割され、各区間に1, 2, ..., hと時計回りに番号が付けられている。そこに ${N}$ 個の建物が存在し、 ${i}$ 番目の建物は区間 ${p_i}$ 内に含まれている。
また、円の中心に立っているChefの視界範囲は ${x}$ であり、円の周囲の連続した ${x}$ 個の区間を一度に見ることができる。

今回、Chefが向く方向をランダムに決定する。
各 ${i = 1, 2, ..., h}$ について ${a_i}$ が決められており、視界に ${i, i+1, ..., i+x}$ (mod ${h}$ )が含まれる確率は ${\frac{a_i}{s}}$ ( ${s = \sum_{i=1}^h a_i}$ )である。

ただし、 ${h}$ がでかいため、 ${a_i}$ は最初の ${K}$ 個しか与えられず、 ${i = K+1, ..., h}$ については係数 ${c_i}$ を用いて以下の式で計算できる。
$a_i = \sum_{j=1}^K c_j a_{i-j}$

ランダムな方向を向いた時のChefの視界の範囲 ${x}$ に含まれる建物の数の期待値をmodulo 163577857で求めなさい。

制約

${1 \le N \le 2000}$
${1 \le h \le 10^9}$
${1 \le K \le 100}$

${1 \le p_i \le h}$
${1 \le x, K \le h}$
${0 \le a_i \le 10^9}$
${0 \le c_i \le 10^9}$

解法

累積和を高速に計算して解いた。

${h}$ は大きい数だが、 ${N}$ が小さいことを考えて、期待値の計算を
$\sum_{i = 1}^h \frac{a_i}{s} \times \text{(番号i, i+1, ..., i+xに含まれる建物数)}$
ではなく、
$\sum_{i = 1}^N \text{(建物iが視界範囲xに含まれる確率)} = \sum_{i = 1}^N \frac{1}{s} \sum_{j = 0}^{x-1} a_{p_i-j}$
(ただし、 ${a_{-i} = a_{h-i} (i \ge 0)}$ とする)
で計算することを考える。

また、累積和 $S_k = \sum_{i = 1}^k a_i$ を用いると、
期待値は $\sum_{i = 1}^N \frac{1}{s} (S_{p_i} - S_{p_i-x})$ で計算できることが分かる。
(ただし、 ${S_{-i} = S_{h-i-1} - S_h (i \ge 0)}$ とする)

あとはこの各 ${S_i}$ を求めればよい。

${S_i}$ の漸化式は ${a_i}$ の漸化式から計算でき、
${S_i = S_{i-1} + a_i = S_{i-1} + \sum_{j=1}^k c_j a_{i-j} = S_{i-1} + \sum_{j=1}^k c_j (S_{i-j} - S_{i-j-1}) \\ = \underline{(c_1 + 1)S_{i-1} + (c_2 - c_1)S_{i-2} + ... + (c_K - c_{K-1})S_{i-K} - c_K S_{i-K-1}}}$
という漸化式が導ける。

この漸化式は線形漸化式であるため、行列を用いた繰り返し二乗法が利用できる。
しかし、 ${S_i}$ を計算するのに ${O(K^3 \log i)}$ となるため、今回の制約上間に合わない。

そこで、繰り返し二乗法よりも高速に ${S_i}$ を計算できるKitamasa法を用いる。
これにより、 ${S_i}$ を ${O(K^2 \log i)}$ で計算できる。

前まとめたKitamasa法メモっぽいの。
smijake3.hatenablog.com

実装

計算量は ${O(K^2 N \log h)}$ 。

提出コード(PyPy2): https://www.codechef.com/viewsolution/18731995

MOD = 163577857
N = int(raw_input())
H = int(raw_input())
X = int(raw_input())
K = int(raw_input())
P = map(int, raw_input().split())
A = map(int, raw_input().split())
C = map(int, raw_input().split())
 
# 累積和: S_1, ..., S_K
S = [0]*(K+1)
for i in xrange(K):
    S[i+1] = S[i] + A[i]
 
# 累積和の漸化式の係数
D = [0]*(K+1)
D[0] = 1
for i in xrange(K):
    D[i] += C[i]
    D[i+1] -= C[i]
    D[i] %= MOD
D[-1] %= MOD
D = D[::-1]
 
# C(n, *) -> C(n+1, *) を計算
def nxt(C):
    C1 = [0]*(K+1)
    C1[0] = C[K] * D[0] % MOD
    for i in xrange(1, K+1):
        C1[i] = (C[i-1] + C[K] * D[i]) % MOD
    return C1
# C(n, *) -> C(2n, *) を計算
def dbl(C):
    Ci = C[:]
    C1 = [0]*(K+1)
    for j in xrange(K+1):
        for i in xrange(K+1):
            C1[i] += C[j] * Ci[i]
        Ci = nxt(Ci)
    for i in xrange(K+1):
        C1[i] %= MOD
    return C1

# Kitamasa法でS_kを計算するための係数c_iを計算
def calc(k):
    bits = []
    while k > K:
        bits.append(k & 1)
        k >>= 1
    C = [0]*(K+1)
    C[k] = 1
    for b in reversed(bits):
        C = dbl(C)
        if b:
            C = nxt(C)
    return C
 
# S_kを計算
def get(k):
    C = calc(k)
    res = 0
    for i in xrange(K+1):
        res += S[i] * C[i]
    return res % MOD
 
# 累積和から期待値を計算
ans = 0
su = get(H)
for p in P:
    if X <= p:
        ans += get(p) - get(p-X)
    else:
        ans += get(p) + (su - get(H - (X-p)))
print(ans * pow(su, MOD-2, MOD) % MOD)