はまやんはまやんはまやん

hamayanhamayan's blog

Coins on the tree [CODE THANKS FESTIVAL 2018 F]

競技プログラミング

https://beta.atcoder.jp/contests/code-thanks-festival-2018-open/tasks/code_thanks_festival_2018_f

前提知識

解説

https://beta.atcoder.jp/contests/code-thanks-festival-2018-open/submissions/3665058

方針は単純で、構築問題テクの『「条件を満たす辞書順最小」頭から貪欲に整合性が保たれるように決めていく』を適用するだけ。
現在コインを置くことができる頂点の中で小さい方から、「その頂点にコインを置いても、その後にコインを置く操作ができるか」をチェックして、できるなら置くということを繰り返す。
　
このチェックのために、2つの関数を用意する。
getmax(m) := 後m枚のコインを置くときに、できる操作の最大回数
getmin(m) := 後m枚のコインを置くときに、できる操作の最小回数
この2つの中身はとても似ている。
ある頂点iにコインを置くときは
① 根から頂点iまでコインが置かれていないと置ける
② 根からの距離+1の操作回数
である。
①は「ok[i] := 頂点iと根の間にコインが無い」というのを持っておけば分かる。
②を考えると、getmaxは②が大きいものからコインを置くと操作回数を最大化できるし、
getminは小さいものからコインを置くと操作回数を最小化できる。
なので、okが適切に管理できれば、getmax, getminは簡単に実装できる。
　
solve関数で実際に貪欲にやっていくが、

rep(i, 0, N) if (ok[i]) {
            vector<int> history;
            dfs2(i, history);
 
            ll mi = getmin(M - dm - 1);
            ll ma = getmax(M - dm - 1);
 
            int dk = dpt[i] + 1;
 
            if (mi <= K - dk and K - dk <= ma) {
                ans.push_back(i);
                ng = 0;
                K -= dk;
                break;
            }
 
            fore(j, history) ok[j] = 1;
        }

頂点iにコインを置くと、その子供が全てok[c] = 0になるので、それをdfs2でやる。
ただ、操作をやめる場合があるので、ok[c]=0にした頂点を全てhistoryで覚えておく。
最後にやっぱり辞めた場合はそれを使ってもとに戻す。
この操作が部分永続の基本的なやり方。
　
あとは、残りのM-dm-1回の操作のmin, maxを求めて、残りの回数K-dkがその中に入っていれば、
後の操作ができるので、採用する。
変数dkは、頂点をiを取るときの操作回数で、(dpt[i]+1)である。
実装がきついのでモジュール毎に分けて実装しよう。

int N, M, K;
vector<int> C[303];
int root;
//---------------------------------------------------------------------------------------------------
int dpt[303];
void dfs(int cu, int d = 0) {
    dpt[cu] = d;
    fore(c, C[cu]) dfs(c, d + 1);
}
//---------------------------------------------------------------------------------------------------
int ok[303];
void dfs2(int cu, vector<int>& history) {
    if (ok[cu]) {
        ok[cu] = 0;
        history.push_back(cu);
    }
    fore(c, C[cu]) dfs2(c, history);
}
//---------------------------------------------------------------------------------------------------
ll getmax(int m) {
    vector<int> v;
    rep(i, 0, N) if (ok[i]) v.push_back(dpt[i]);
    sort(all(v), greater<int>());
    if (v.size() < m) return -infl;
 
    ll res = 0;
    rep(i, 0, m) res += v[i] + 1;
    return res;
}
ll getmin(int m) {
    vector<int> v;
    rep(i, 0, N) if (ok[i]) v.push_back(dpt[i]);
    sort(all(v));
    if (v.size() < m) return infl;
 
    ll res = 0;
    rep(i, 0, m) res += v[i] + 1;
    return res;
}
//---------------------------------------------------------------------------------------------------
vector<int> solve() {
    dfs(root);
    rep(i, 0, N) ok[i] = 1;
 
    // -1のケース
    if (K < getmin(M) or getmax(M) < K) return {};
 
    vector<int> ans;
    rep(dm, 0, M) {
        int ng = 1;
        rep(i, 0, N) if (ok[i]) {
            vector<int> history;
            dfs2(i, history);
 
            ll mi = getmin(M - dm - 1);
            ll ma = getmax(M - dm - 1);
 
            int dk = dpt[i] + 1;
 
            if (mi <= K - dk and K - dk <= ma) {
                ans.push_back(i);
                ng = 0;
                K -= dk;
                break;
            }
 
            fore(j, history) ok[j] = 1;
        }
        if (ng) return {};
    }
 
    return ans;
}
//---------------------------------------------------------------------------------------------------
void _main() {
    cin >> N >> M >> K;
 
    rep(i, 0, N) {
        int P; cin >> P; P--;
        if (P < 0) root = i;
        else {
            C[P].push_back(i);
        }
    }
 
    auto ans = solve();
 
    if (ans.size() == 0) printf("-1\n");
    else {
        rep(i, 0, M) {
            if (i) printf(" ");
            printf("%d", ans[i] + 1);
        }
        printf("\n");
    }
}