けんちょんの競プロ精進記録

競プロの精進記録や小ネタを書いていきます

AtCoder AGC 010 B - Boxes (500 点)

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

問題へのリンク

問題概要

 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;
}