【AtCoder】ABC 388 D - Coming of Age Celebration | 緑コーダーが解くAtCoder

atcoder.jp

配点: 400 点 / 実行時間制限: 2 sec / メモリ制限: 1024 MB / Difficulty: 659

成人祝い...宇宙人...星人... うーん、妙だな。*1 *2

問題概要
制約
考察
実装例
コメント

問題概要

ある星には $N$ 人の宇宙人がおり、全員未成年である。 $i$ 人目の宇宙人は現在 $A_i$ 個の石を所持しており、ちょうど $i$ 年後に成人する。

この星では誰かが成人するとき、石を1個以上所持している成人全員が、成人する宇宙人に成人祝いとして石を1個渡す。このとき、 $N$ 年後に各宇宙人が所持している石の個数を求めよ。ただし、今後新たな宇宙人は産まれないものとする。

制約

$1 \leq N \leq 5 \times 10^5$
$0 \leq A_i \leq 5 \times 10^5$
入力はすべて整数。

考察

簡単のため、以下では $i$ 人目の宇宙人を単に「宇宙人 $i$ 」と表記する。

まず、問題文の読み替えを行おう。成人時に「自身より年下(左側)の宇宙人から石をもらう」のではなくて、「自身より年上(右側)の宇宙人に対して可能な限り渡す」と考える。つまり、

宇宙人 $i$ が成人したとき、宇宙人 $i+1, i+2, \cdots$ へ自分の石がなくなるまで石を渡していく

ということ。こうすることで、石の移動を区間で考えやすくなったことが分かるだろうか。

宇宙人 $i$ が成人時に宇宙人 $i+1, i+2, \cdots, i+1+\mathrm{give\_stones}_i$ に石を渡すとする。ここで、 $\mathrm{give\_stones}_i = \min(A_i, N - i)$ である。

石を各々に1個ずつ渡す操作は「 $A_j \: (i+1 \leq j \leq i+1+\mathrm{give\_stones}_i)$ に1加算する」というようにしてもいいのだが、これでは全体の計算量が $O(N^2)$ となってしまい、 TLE になる。

そこで、 imos 法により高速化を試みる。具体的には以下の通り。なお、ここからは 0-indexed での説明となる。

$\mathrm{get\_stones}_{i} \: (0 \leq i \leq N - 1) :=$ 宇宙人 $i$ が受け取る石の合計
$\mathrm{diff}_{i} \: (0 \leq i \leq N) :=$ imos 法のテーブル

を用意し、 $i = 0, 1, \cdots, N - 1$ の順に次の処理を行う。

のとき(宇宙人は石を渡されることはない)、
- $\mathrm{get\_stones}_{i} = \mathrm{get\_stones}_{i-1} + \mathrm{diff}_{i}$ として( $\mathrm{diff}_{i}$ までの累積和を取った)、 $A_i$ に加算する。 ※自身より左の宇宙人が受け取れる石の数はこれ以上変化しないので、累積和を逐次取ってしまって構わない。
$\mathrm{give\_stones}_i = \min(A_i, N - 1- i)$ として、 $A_i$ から減算する。
$\mathrm{diff}_{i+1}$ をインクリメント、 $\mathrm{diff}_{\min(N, i+1+\mathrm{give\_stones}_i)}$ をデクリメントする(imos 法テーブルの更新)。

こうすることで、全体の計算量を $O(2N)$ (定数倍を無視すると $O(N)$ )にまで落とすことができた。

実装例

#include <bits/stdc++.h>
using namespace std;

#define rep(i, start, end) for (auto i = (start); (i) < (end); (i)++)

// ======================================== //

int main()
{
    int N;
    cin >> N;
    vector<int> A(N);
    rep(i, 0, N) cin >> A[i];
    vector<int> diff(N + 1, 0), get_stones(N, 0);
    rep(i, 0, N)
    {
        if (i > 0)
        {
            get_stones[i] = get_stones[i - 1] + diff[i];
            A[i] += get_stones[i];
        }

        int give_stones = min(A[i], N - i - 1);
        A[i] -= give_stones;

        diff[i + 1]++;
        diff[min(N, i + 1 + give_stones)]--;
    }

    rep(i, 0, N)
    {
        cout << A[i] << " ";
    }
    cout << endl;

    return 0;
}