Official

D - Coming of Age Celebration Editorial by en_translator


We simply call the \(i\)-th alien “alien \(i\).”

By focusing not on receiving but on giving stones, we can find the correct answer by interpreting the problem as follows:

Alien \(i\) has \(A_i\) stones and will become an adult \(i\) years later.

Alien \(i\) gives each of aliens \(i + 1, i + 2, \ldots, N\) a stone (if he has one) in this order when he becomes an adult.

How many stones will they finally have?

From now on, we consider the rephrased problem.

Prepare a sequence \(C\) of length \(N\) and a sequence \(D\) of length \((N+1)\), initially filled with \(0\). We want \(C_i\) eventually be the total number of stones that alien \(i\) will receive.

Suppose that alien \(i\) will give aliens \(i + 1, i + 2, \ldots, R_i\) after becoming an adult. Here, we want to add \(1\) to \(C_j\) for each \(j = i + 1, i + 2, \ldots, R_i\), but naively doing so costs a total time complexity of \(\Theta(N^2)\). Instead, we use the cumulative sum trick to optimize it. \(D\) serves as a differential sequence of \(C\).

  • For each \(i = 1, 2, \ldots, N\), do the following.
    • (If \(i > 1\)) let \(C_i\) be \(C_{i - 1} + D_i\). This determines the number of stones when alien \(i\) becomes an adult, yielding \(R_i\).
    • Add \(D_{i + 1}\) to \(1\) and \(D_{R_i + 1}\) to \(-1\).

This procedure can be done in a total of \(O(N)\) time.

Sample code (0-based indexing)

#include <bits/stdc++.h>
#define rep(i, n) for (int i = 0; i < (n); i++)
using namespace std;

int main() {
	int n;
	cin >> n;
	vector<int> a(n);
	rep(i, n) cin >> a[i];
	vector<int> c(n), d(n + 1);
	rep(i, n) {
		if (i != 0) {
			c[i] = c[i - 1] + d[i];
			a[i] += c[i];
		}
		int cnt = min(n - i - 1, a[i]);
		a[i] -= cnt;
		d[i + 1]++;
		d[min(n, i + cnt + 1)]--;
	}
	rep(i, n) cout << a[i] << " \n"[i == n - 1];
}

posted:
last update: