日々drdrする人のメモ

今日も明日もdrdr

CodeChef June Challenge 2018: Expected Buildings

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}を用いて以下の式で計算できる。
 {\displaystyle 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}が小さいことを考えて、期待値の計算を
 {\displaystyle \sum_{i = 1}^h \frac{a_i}{s} \times \text{(番号i, i+1, ..., i+xに含まれる建物数)}}
ではなく、
 {\displaystyle \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)}とする)
で計算することを考える。

また、累積和 {\displaystyle S_k = \sum_{i = 1}^k a_i }を用いると、
期待値は  {\displaystyle \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)