【AtCoder】ABC 480 D - Concatenation of Prefixes

D - Concatenation of Prefixes

AtCoder is a programming contest site for anyone from beginners to experts. We hold weekly programming contests online.

AtCoder

実行時間制限: 2 sec / メモリ制限: 1024 MiB / Difficulty: 864 / NoviSteps: 1Q / 配点: 400 点

問題概要

長さ $N$ の整数列 $(A_1,\dots,A_N)$ が与えられる。

以下の手順で長さ $\frac{N(N+1)}{2}$ の数列 $B$ を作る。

• 最初、 $B$ を空列とする。
• $i=1,2,\dots,N$ の順に、$B$ の末尾に $(A_1,\dots,A_i)$ を連結する。

以下の形式のクエリが $Q$ 個与えられるので、それぞれ処理せよ。

• $L,R$ が与えられる。$B_L+B_{L+1}+\dots+B_R$ を求めよ。

制約

• $1 \leq N \leq 10^5$
• $1 \leq Q \leq 10^5$
• $-10^8 \leq A_i \leq 10^8$
• 各クエリにおいて、$1\leq L\leq R \leq \frac{N(N+1)}{2}$
• 入力は全て整数

考察

$G_i = (A_1, A_2, \dots, A_i)$ とすると、 $B = G_1 + G_2 + \dots + G_N$ である。

このとき、$G_i$ の末尾の項は $B$ の $\frac{i(i+1)}{2}$ 番目に現れる。 これを $C_i$ とおく。

いま、$f(x) = \sum_{i=1}^{x} B_i$ とおけば($B$ の累積和)、各クエリでは $f(R) - f(L-1)$ を計算すればよい。

$x$ が $G_k$ に含まれる要素であるとすると、 $k = \min\{i \mid C_i \geq x\}$ と定義できるから、これは二分探索で求めることができる。

すると、$G_k$ より前にある要素の個数は $C_{k-1}$ 個なので、$G_k$ の中での位置は $x - C_{k-1}$ と表せる。

したがって、

• $S_i = \sum_{j=1}^{i} A_j$ ($A$ の累積和)
• $T_i = \sum_{j=1}^{i} S_j$ ($S$ の累積和)

を定義すれば、

と表すことができる。

計算量は、累積和の計算に $O(N)$、各クエリに対して二分探索を行うので $O(Q\log N)$ となり、全体で $O(N + Q\log N)$ となる。

実装例

1.#include <bits/stdc++.h>
2.using namespace std;
3. 
4.#define all(x) (x).begin(), (x).end()
5.#define rep(i, start, end) for (auto i = (start); (i) < (end); (i)++)
6. 
7.using ll = long long;
8. 
9.// ======================================== //
10. 
11.int main()
12.{
13.    int N, Q;
14.    cin >> N >> Q;
15.    vector<ll> A(N);
16.    rep(i, 0, N) cin >> A[i];
17. 
18.    vector<ll> S(N + 1, 0), T(N + 1, 0), C(N + 1, 0);
19.    rep(i, 0, N)
20.    {
21.        S[i + 1] = S[i] + A[i];
22.        T[i + 1] = T[i] + S[i + 1];
23.        C[i + 1] = C[i] + (i + 1);
24.    }
25. 
26.    auto f = [&](ll x) -> ll {
27.        if (x == 0) return 0;
28. 
29.        int k = lower_bound(all(C), x) - C.begin();
30.        return T[k - 1] + S[x - C[k - 1]];
31.        };
32. 
33.    while (Q--)
34.    {
35.        ll L, R;
36.        cin >> L >> R;
37. 
38.        cout << f(R) - f(L - 1) << endl;
39.    }
40. 
41.    return 0;
42.}

https://atcoder.jp/contests/abc480/submissions/80170230

atcoder.jp

実装時間: 30分

コメント

累積和が大量に出てくるので、気を付けないと添え字ガチャに苦しむことになる。