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

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

AtCoder AGC 008 D - K-th K (800 点)

実装は手こずったものの、それなりに自信のある Greedy を提出できて一発 AC できてよかった

問題へのリンク

問題概要

1 〜  N の値を  N 個ずつもつ長さ  N^{2} の数列であって、

  •  i について、 x_i 番目の値が  i 個目の  i である

という条件を満たすものを 1 つ構築せよ。条件を満たすものが存在しない場合には -1 をリターンせよ。

制約

  •  1 \le N \le 500

考えたこと

面白そうな構築系!!!!!!!

この問題、まず 1 を埋めて、2 を埋めて... という風に考えたくなるのだが、そのような思考だと、2 を先に埋めるよりも 4 が割と先頭の方に条件として指定されていて、4 を先に使わないと「4 が来る前に 4 を 3 個消費しないといけないのにそれができずに詰む」ということになったりする。

なので、思考を変えて、


  •  x_i 番目までに  i i-1 個消費しないといけない

という条件を  x_i の値が小さい方から優先的に消化していく。

そして消化しきってしまっても、まだ次の  x_i までマス目が余っているならば、既に通過した  i に関してはその後ガンガン使って良い (全部使ってしまったら No)


という風にすればよさそう。なぜこの Greedy でよいかだが、答えがもし "Yes" だとして条件を満たす数列が得られていたならば、必ず上記の優先順位を満たすように数値を入れ替えていくことができる。

さて反対に、この Greedy をやっても、

  1.  x_i 番目までに  i i-1 個消費できなかった (N = 2 で「先頭に 2」が課せられるなど)
  2. 未来訪れる  x_i について、全部ちゃんと  i-1 個消費しきっており、しかも既に  x_i を通過した各  i についても全部  N 個使い切ってしまった

という状況になりえて、そういうケースではこの Greedy に限らずいかなる構築を試みようとも明らかに "No" になるし、こういう風にさえならなければ必ず上記の Greedy によって解を構築できる。

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

int N;
vector<int> x;

using pint = pair<int,int>;
vector<int> po; // その場所に制約があるなら、そこに置かれる数、なければ 0
vector<int> con; // con[i] := i 番目の checkpoint についての数

int main() {
    int N; cin >> N;
    x.resize(N); for (int i = 0; i < N; ++i) cin >> x[i], --x[i];
    po.resize(N*N, 0);
    for (int i = 0; i < N; ++i) po[x[i]] = i + 1;
    con.clear();
    for (int i = 0; i < N*N; ++i) if (po[i]) con.push_back(po[i]);
    
    int iter = 0;
    vector<int> num(N+2, 0); // num[i] := 数 i を何個使ったか
    vector<int> res(N*N, 0);
    set<int> alr;
    bool ok = true;
    for (int i = 0; i < N*N; ++i) {
        if (!ok) break;
        if (po[i]) {
            if (num[po[i]] < po[i]-1) { ok = false; break; }
            res[i] = po[i];
            ++num[po[i]];
            alr.insert(po[i]);
        }
        else {
            while (iter < N && num[con[iter]] >= con[iter]-1) ++iter;
            if (iter == N) {
                int canuse = 1;
                while (canuse <= N && (!alr.count(canuse) || num[canuse] == N)) ++canuse;
                if (canuse == N + 1) { ok = false; break; }
                res[i] = canuse;
                ++num[canuse];
            }
            else {
                res[i] = con[iter];
                ++num[con[iter]];
            }
        }
    }
    
    if (!ok) puts("No");
    else {
        puts("Yes");
        for (int i = 0; i < N*N; ++i) cout << res[i] << " ";
        cout << endl;
    }
}