Coins on the tree

問題ページ - F - Coins on the tree

概要

頂点の根付き木が与えられる。頂点の親はであり、である頂点が根である。この根付き木上で、枚のコインを用いて以下のいずれかの操作を合わせてちょうど回行う。

• 根にコインが置かれていない場合、新しいコインを根に置く
• 既にコインが置かれている頂点を1つ選び、コインが置かれていない子のいずれかにコインを移動する

番目に置いたコインが置かれている頂点の番号をとするとき、列のうちの辞書順最小を求めよ。ただし、枚のコインを全て置くことができない場合や、ちょうど回の操作を行うことができない場合はを出力せよ。

制約

• 
• 
• 
• 与えられるグラフは根付き木である
• 入力はすべて整数である

解法

辞書順最小を求めたいので、前から貪欲に決めていく。ここで、枚目のコインを頂点に置くことが可能であることの条件は、以下の通りである。

• 頂点は枚目のコインを置いた頂点を根とする部分木に含まれない
• 枚目のコインを置いた頂点を根とする部分木を元の木から取り除いたとき、頂点数が個以上である
• 枚目のコインを置く際に行った操作回数の和をとする。このとき、は枚目のコインを置く際に行う操作回数の和としてありえるもののうちの最小以上かつ最大以下である

1つ目の条件は、各頂点にコインが置かれているか否かをbool型の配列で保持しておくことで、DFSを用いて根方向にたどっていくことで時間で求められる。2つ目の条件と3つ目の条件については、これらを求めるために深さの頂点がいくつあるかを配列で保持しておく。ただし、ここでは根の深さをとする。DFSで子孫方向にたどっていき、各深さの頂点数を1ずつ減らしていくことで部分木を取り除く処理と同じことをできる。ただし、重複して取り除くことを避けるため、コインが置かれている頂点に着いた場合は探索を打ち切る。こうすることで、残りの頂点数は深さの頂点の数の和となるため、2つ目の条件も時間で求められる。3つ目の条件は、これを満たしていれば操作回数はあとで辻褄を合わせることができる。この条件は、深さが小さい方から個の頂点の深さの和が操作回数の最小となり、深さが大きい方から個の頂点の深さの和が操作回数の最大となることから、これも時間で求められる。これらの条件を満たさなかった場合は、取り除いた部分木を元に戻す必要があることに注意する。枚のコインそれぞれについて個の頂点を前から順に使えるかどうかを調べるため、全体の計算量はとなる。

提出コード

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

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

const int IINF = 1e9 + 100;

int n, m, k, root;
vector<int> p, dis, dis_cnt, ans;
vector< vector<int> > g;
vector<bool> used;

void calc_depth(int v, int depth){
    dis[v] = depth;
    ++dis_cnt[depth];
    for(auto x : g[v]){
        if(dis[x] > depth + 1){
            calc_depth(x, depth + 1);
        }
    }
    return;
}

bool ancestor(int v){
    if(used[v]){
        return false;
    }
    if(v == root){
        return true;
    }
    return ancestor(p[v]);
}

void descendant(int v, bool mode){
    if(used[v]){
        return;
    }
    if(mode){
        --dis_cnt[dis[v]];
    }else{
        ++dis_cnt[dis[v]];
    }
    for(auto x : g[v]){
        descendant(x, mode);
    }
    return;
}

bool check(int node, int num, int ope){
    if(used[node] || ope + dis[node] > k || !ancestor(node)){
        return false;
    }

    descendant(node, true);
    
    int node_cnt = 0;
    for(int i = 0 ; i <= n ; ++i){
        node_cnt += dis_cnt[i];
    }
    if(node_cnt < m - num - 1){
        descendant(node, false);
        return false;
    }

    int ope_min = 0, ope_max = 0;
    node_cnt = m - num - 1;
    for(int i = 0 ; i <= n ; ++i){
        ope_min += i * min(dis_cnt[i], max(node_cnt, 0));
        node_cnt = max(node_cnt - dis_cnt[i], 0);
    }
    node_cnt = m - num - 1;
    for(int i = n ; i >= 0 ; --i){
        ope_max += i * min(dis_cnt[i], max(node_cnt, 0));
        node_cnt = max(node_cnt - dis_cnt[i], 0);
    }

    if(ope_min <= k - ope - dis[node] && k - ope - dis[node] <= ope_max){
        return true;
    }else{
        descendant(node, false);
        return false;
    }
}

bool solve(int num, int ope){
    if(num == m){
        return ope == k;
    }
    for(int i = 0 ; i < n ; ++i){
        if(check(i, num, ope)){
            ans.push_back(i);
            used[i] = true;
            if(solve(num + 1, ope + dis[i])){
                return true;
            }
            ans.pop_back();
            used[i] = false;
            descendant(i, false);
        }
    }
    return false;
}

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

    cin >> n >> m >> k;

    p.resize(n);
    dis.resize(n, IINF);
    dis_cnt.resize(n + 1);
    g.resize(n);
    used.resize(n);

    for(int i = 0 ; i < n ; ++i){
        cin >> p[i];
        if(p[i] == 0){
            root = i;
            continue;
        }
        --p[i];
        g[p[i]].push_back(i);
    }

    calc_depth(root, 1);

    if(solve(0, 0)){
        for(int i = 0 ; i < m ; ++i){
            cout << ans[i] + 1;
            if(i < m - 1){
                cout << " ";
            }else{
                cout << endl;
            }
        }
    }else{
        cout << -1 << endl;
    }

    return 0;
}