こどふぉ特有の $X$ 回以下の操作で〜を達成せよ、という問題

問題へのリンク

問題概要

長さ $N$ の正の整数からなる数列 $A_{1}, \dots, A_{N}$ が与えられる。以下の操作を $3N$ 回以下繰り返すことで、数列の値がすべて等しくなるようにしたい。そのような操作列を一つ求めよ。不可能である場合は -1 を出力せよ。

[操作]：整数 $i, j, x$ を選び、 $A_{i}$ の値を $i \times x$ だけ減少させて、 $A_{j}$ の値を $i \times x$ だけ増加させる。ただし $A_{i}$ の値が負になってはいけない

制約

$1 \le N \le 10^{4}$

考えたこと

まず、操作をどのように行っても、 $A$ の総和が不変であることに注目する。よって、明らかに $A$ の総和 $S$ が $N$ の倍数でなければならない。

このとき $V = \frac{S}{N}$ とする。すべての要素をうまいこと $V$ に揃えたい。色々迷走してしまったが、最終的にたどり着いたのは、次の方針だ。

フェーズ I： $A_{1}$ 以外のすべての要素を 0 にする (それが可能であることを後に示す)
フェーズ II： $A_{1}$ の要素を各要素に分配していって実現する

フェーズ I が実現可能ならば、フェーズ II は簡単だ。単に $A_{1}$ から各 $A_{i}$ に対して $x = V$ として操作を行なっていけばよい。

フェーズ I は、実は $i = 2, 3, \dots, N$ の順序で行なっていけばよいのだ。最悪ケースは

$A = (1, 1, 1, 1, 1, 4, 5)$

のように、ひたすら 1 が並ぶケース。しかし $A_{i} \ge 1$ であることに注目しよう。そうすると、

$A_{1} + A_{2} \ge 2$

であることが保証される。よって、一旦 $A_{1}$ から $A_{2}$ へ要素を移して $A = (0, 2, ...)$ の状態にしてから、 $A_{2}$ から $A_{1}$ に戻せば $A = (2, 0, \dots)$ の状態になる。

より一般に、 $i = 2, 3, \dots, N$ の順に

$A_{i}$ が $i$ の倍数になるように、 $A_{1}$ から適切な量を $A_{i}$ に配る
$A_{i}$ が 0 になるように $A_{i}$ から $A_{1}$ へ配る

とすれば OK。これで最大でも $3(N-1)$ 回の操作で実現できる。

コード

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

void solve(vector<long long> &A) {
    int N = A.size();
    long long sum = 0;
    for (auto v: A) sum += v;
    if (sum % N != 0) {
        cout << -1 << endl;
        return;
    }
    long long num = sum / N;

    // 0-indexed
    vector<long long> ri, rj, rx;
    auto run = [&](int i, int j, long long x) -> void {
        //cout << i+1 << " " << j+1 << " " << x << endl;
        if (x == 0) return;
        ri.push_back(i+1), rj.push_back(j+1), rx.push_back(x);
        A[i] -= x * (i+1);
        A[j] += x * (i+1);
    };

    // 0 に集める
    for (int i = 1; i < N; ++i) {
        long long p = A[i] / (i+1), r = A[i] % (i+1);
        if (r == 0) run(i, 0, p);
        else {
            run(0, i, i+1-r);
            run(i, 0, A[i]/(i+1));
        }
    }

    // 0 から揃えていく
    for (int i = 1; i < N; ++i) {
        long long dif = num - A[i];
        run(0, i, dif);
    }

    // 出力
    cout << ri.size() << endl;
    for (int i = 0; i < ri.size(); ++i) {
        cout << ri[i] << " " << rj[i] << " " << rx[i] << endl;
    }
}

int main() {
    cin.tie(0); 
    ios::sync_with_stdio(false);

    int T;
    cin >> T;
    while (T--) {
        int N;
        cin >> N;
        vector<long long> A(N);
        for (int i = 0; i < N; ++i) cin >> A[i];
        solve(A);
    }
}

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

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

Codeforces Round #673 (Div. 1) B. Make Them Equal (R2000)

問題概要

制約

考えたこと

コード