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

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

AtCoder ABC 379 C - Sowing Stones (1Q, 緑色, 300 点)

AtCoder AtCoder300点 ABC-C 緑色diff NoviSteps1Q データ構造テク：indexベースで考える データ構造 必要条件を列挙したら十分条件になる Greedy 条件の言い換え 操作 最適化問題 最小コスト 最小回数・最小個数を求める 一次元グリッド

これはややこしい！

問題概要

マス 1, 2, \dots, N が一列に並んでいる。最初、 M 個のマスに石が入っており、マス X_{i} には A_{i} 個の石が入っている。あなたは以下の操作を好きな回数行うことができる。

「マス i に石があるとき、マス i からマス i+1 へと石を 1 つ移動させる」

N 個のマスすべてに石がちょうど 1 個ずつ入っている状態にするために必要な操作回数の最小値を求めよ（不可能な場合は -1 と答えよ）。

制約

  • 2 \le N \le 2 \times 10^{9}
  • 1 \le M \le 2 \times 10^{5}

考えたこと

まず自明にわかることは「 A_{i} の総和 ≠ M のときは不可能」ということだ。よって、以降、 A_{i} の総和 = M であるとしよう。

さて、石の分布を表すのに「どのマスに何個の石があるか」ではなく、「各石がどのマスにあるか」で考えることにしよう。具体的には、マスの左側にある石から順に、その石のあるマスの番号を並べたベクトルを考えることにする。たとえば、サンプル 1 の初期状態（マス 1 に 3 個、マス 4 に 2 個）は

(1, 1, 1, 4, 4)

のように表せる。これを

(1, 2, 3, 4, 5)

の状態にするのが目標だ。なお、このベクトルによる表現方法を用いると、問題文の操作は「ベクトルの 1 要素を選び、値を +1 する（さらに小さい順に並び替える）」と言い換えることができる。よって、もとの問題は次のように言い換えられる。

N 個の値からなる数列 B_{1}, B_{2}, \dots, B_{N} が与えられる（この数列はランレングス圧縮された状態で与えられ、 M 種類の値からなる）。この数列に対して

  • 数列の要素を 1 つ選び、1 を足して、
  • 数列全体を小さい順に並び替える

という操作を最小回数行うことで、数列が (1, 2, \dots, N) となるようにせよ。

言い換えた問題の考察

まず、もし

B_{i} \gt i

となるような i があったならば、不可能であることが言える。なぜならば、上記の操作によって、どの i についても B_{i} が減少することは決してないからだ。

よって、今後は任意の i に対して、 B_{i} \le i であるとしよう。この場合は、必ず可能である。具体的には、 B_{N}, B_{N-1}, \dots, B_{1} の順に、 N, N-1, \dots, 1 になるまで値を増やし続ければよいのだ。

この操作に要する回数は、

\sum_{i=1}^{N} i - \sum_{i=1}^{N} B_{i}
= \frac{N(N+1)}{2} - \sum_{i=1}^{M} A_{i}X_{i}

となる。そして実は、数列 B を (1, 2, \dots, N) にすることができる操作列はさまざまなものが考えられるが、そのいずれも操作回数は上の式になる。なぜならば、操作によって B の総和はちょうど 1 ずつ増えるからだ。

まとめ

以上をまとめると、次のようになる。

  • \sum_{i=1}^{M} A_{i} \neq N のとき：不可能
  • 上にように定義した数列 B について、 B_{i} \gt i となる i が存在するとき：不可能
  • それ以外のとき：最小の捜査回数は \frac{N(N+1)}{2} - \sum_{i=1}^{M} A_{i}X_{i}

この解法は O(M) の計算量で実装できる。

コード

ここではマス目を 0-indexed で実装した。

#include <bits/stdc++.h>
using namespace std;
using pll = pair<long long, long long>;

int main() {
    long long N, M, sum = 0;
    cin >> N >> M;
    vector<pll> stone(M);
    for (int i = 0; i < M; i++) cin >> stone[i].first, stone[i].first--;
    for (int i = 0; i < M; i++) cin >> stone[i].second, sum += stone[i].second;

    if (sum != N) {
        cout << -1 << endl;
        return 0;
    }

    sort(stone.begin(), stone.end());
    long long index = 0, ax_sum = 0;
    for (auto [place, num] : stone) {
        if (place > index) {
            cout << -1 << endl;
            return 0;
        }
        index += num;
        ax_sum += place * num;
    }
    cout << N * (N - 1) / 2 - ax_sum << endl;
}