連続する部分列の和の総和
Table of contents
問題
- 長さ $N$ の数列 $A = [a_1, a_2, ..., a_N]$ が与えられる
- この数列の任意の連続する部分列の和の総和を $10^9+7$ で割った余りを求めよ
例
$A = [1, 2, 3]$ のとき、和は $1+2+3+(1+2)+(2+3)+(1+2+3)=20$
制約
- $1 \le N \le 10^6$
- $-10^9 \le a_i \le 10^9 (i = 1, ..., N)$
解法
- $A$ の累積和を取った数列を $S$ とする
- つまり、$S_i := \sum_{j=1}^{i}a_j$
- ただし、$S_0 = 0$ とする
- このとき、求める総和は、$\sum_{(i,j):0 \le i < j \le N}(S_j - S_i)$ と書ける
- $(i,j):0 \le i < j \le N$ は $0 \le i < j \le N$ を満たす組 $(i,j)$ を意味する
- これは和と差しか出てこないため、正の項 $S_j$ と負の項 $-S_i$ がそれぞれ何回この式の中で現れるかを考えればよい
- 正の項 $S_j(j=0,1,...,N)$ が何回現れるかを考えると、$S_0$ は0回、$S_1$ は 1 回 ($i=0$ のとき)、$S_2$ は 2 回 ($i=0,1$ のとき)、と考えていくと、正の項 $S_j$ は $j$ 回現れる
- 負の項 $S_i(i=0,...,N-1)$ が何回現れるかを考えると、$S_0$ は $N$ 回 ($j=1,2,...,N$ のとき)、$S_1$ は $N-1$ 回 ($j=2,i...,N$ のとき) と考えていくと、$S_i$ は $(N-i)$ 回現れる
- よって、$\sum_{(i,j):0 \le i < j \le N}(S_j - S_i)=\sum_{i=0}^{N}(i \times S_i - (N-i)\times S_i)=\sum_{i=0}^{N}(2i-N)S_i$ となる
- 以上より、$O(N)$ で問題を解くことができた
コード
def sum_subarrays(A, mod=10**9+7):
N = len(A)
S = [0]
for a in A:
S.append(S[-1] + a)
S[-1] %= mod
return sum((2 * i - N) * s % mod for i, s in enumerate(S)) % mod
- 負の剰余が出てくるが、Python (2/3系ともに) では負の数 $x(x>0)$ の正の数 $m(m>0)$ についての剰余 $x\%m$は$(x+m)\%m$ と一致するため、問題ない
- C++ ではこのコードは問題が出る
- 負の数の剰余の計算で負の数が出る場合の計算は、言語仕様によって異なるため、注意したほうがよい
類題
- 連続でなくてよい場合は、総和において各数は $2^{N-1}$ 回出現する
- よって答えは $2^{N-1}\times\sum_{i=1}^{N}a_i$
- 和ではなく積の場合も同様に求められる
- $P_0 := 1, P_i := \prod_{j=1}^{i}a_j (i=1,\ldots,N)$ とすると、$\prod_{(i,j):0 \le i < j \le N} \frac{P_j}{P_i}=\prod_{i=0}^{N} \frac{P_i^i}{P_i^{N-i}}=\prod_{i=0}^{N}P_i^{2i-N}$
- 和ではなく bitwise xor の場合も同様で、ビットごとに偶数回か奇数回かを考えればよい
- 連続する部分列の長さが $K$ 以下という制約がある場合
- これは正の項 $S_i$、負の項 $-S_i$ が出現する回数が $K$ 回以下という条件だと考えればよいので、答えは $\sum_{i=0}^{N}(\min(i, K) \times S_i - \min(N-i, K)\times S_i)$
- 連続する部分列の長さが $K$ 以上という制約がある場合
- $\sum_{i=0}^{N}(\max(i-K+1, 0) \times S_i - \max(N-i-K+1, 0)\times S_i)$