状況がゴチャゴチャとして来て、整理するのが大変だった

問題概要

$N$ 個のマスがあって最初は全マスに $0$ が書かれている。これに以下の操作を好きな回数だけ好きな順序で行って数列 $a_1, a_2, \dots, a_N$ にできるかどうか判定せよ。

$(1, 2, 3, \dots, N)$ を巡回させてできるものを足す

制約

$1 \le N \le 10^{5}$
$1 \le A_i \le 10^{9}$

考えたこと

$a$ の中で最大値に着目したとき、それが一度も $N$ を足さずに最大にはなれないことは容易にわかる。

なぜなら、 $N$ 以外のものを足す時、その隣に要素には常に「1 大きい数」が足されていくことになるからだ。よってその場合には最大の数は最大にはなりえない。

ゆえに、「最大値に $N$ を足した瞬間があった」ということが言えるので、

数列 $a$ に対し、最大値から $N$ を引くような操作 (最大値から巡って順に $N, 1, 2, \dots, N-1$ を引くような操作) をずっと行った結果が $0, 0, \dots, 0$ になるかを判定する

ことで解くことができることがわかった。しかし毎回数列から等差数列を引く操作を高速にシミュレートできなければ、TLE してしまう。

操作の高速化 - 数列の隣接差をとる

区間に等差数列を足し引きする処理はセグメントツリーなどで無理矢理やってもよいが、もっと楽にやりたい。等差数列加算操作では、隣接した要素の差をとるとほとんど 1 になることに着目するとよいイメージがある。

今回、

元の数列に対し、 $b_{i} = a_{i+1} - a_{i}$ ( $i+1 = N$ のときは $i+1 = 0$ に) とした世界で考える
そうすると、操作は $(1, 1, -N+1, 1, 1)$ のような 1 箇所だけ $-N+1$ で残りが $1$ であるベクトルを足す操作になる

ということがわかるのでさらに、操作回数を $k$ とすると (操作回数は「総和 / ( $1 + 2 + \dots + N$ 」で求められる)、

元の数列に対し、 $b_{i} = a_{i+1} - a_{i} - k$ とした世界で考える
そうすると、操作は $(0, 0, -N, 0, 0)$ のような 1 箇所だけ $-N$ を足す操作になる

ということがわかる。よって

各 $b$ に対して、「それが $-N$ の倍数で、かつ負の値である」かどうかを判定する

ことで問題を解くことができる。

#include <iostream>
#include <vector>
using namespace std;

bool solve(long long N, const vector<long long> &a) {
    long long S = 0;
    for (int i = 0; i < N; ++i) S += a[i];
    if (S % (N * (N+1)/2) != 0) return false;
    long long num = S / (N * (N+1)/2);
    for (int i = 0; i < N; ++i) {
        long long b = a[(i+1)%N] - a[i] - num;
        if (b % N != 0 || b > 0) return false;
    }
    return true;
}

int main() {
    long long N; cin >> N;
    vector<long long> a(N);
    for (int i = 0; i < N; ++i) cin >> a[i];
    if (solve(N, a)) cout << "YES" << endl;
    else cout << "NO" << endl;
}