The basic idea here is the same as given in the Editorial. So please read that first. Now to efficiently do this in O(n) we need to keep track of possible jumps Pekora can make from a given index.
Suppose we have an initial array of {4, 2, 2, 2, 3, 2}. We will traverse from left to right. Now we will start from index 0 (with value 4) because it is greater than 1. We will take it as the starting point. One thing we can observe is that we need to use it 3 times until it becomes 1. In doing so we will make a jump to 4th, 3rd and 2nd index in the 1st, 2nd and 3rd pass respectively. So all we have to do is find a way to store this so that when we come to a new index we already know how many times we will come to this index. I used an array for this to store the range which will be affected by the current index.
My Submission: 108827845
auto solve(vector<int> &A) {
int i, j, diff, n = A.size();
ll ans = 0;
vector<int> C(n + 1);
auto add = [&](int i, int x) {
if (i > n) i = n;
C[i] += x;
};
auto sub = [&](int i, int x) {
if (i > n) i = n;
C[i] -= x;
};
for (i = 0; i < n; i++) {
if (i) C[i] += C[i - 1];
diff = A[i] - max(A[i] - C[i], 1);
add(i + A[i] - diff + 1, 1);
sub(i + A[i] + 1, 1);
A[i] -= diff;
if (C[i] > diff) {
add(i + 1, C[i] - diff);
sub(i + 2, C[i] - diff);
}
if (A[i] > 1) {
add(i + 2, 1);
sub(i + A[i] + 1, 1);
ans += A[i] - 1;
A[i] = 1;
}
}
return ans;
}