제출 #1320078

#	제출 시각	아이디	문제	언어	결과	실행 시간	메모리
1320078		turral	Railway (BOI17_railway)	C++20	100 / 100	148 ms	52776 KiB

railway

#include <iostream>
#include <vector>
#include <map>

using namespace std;

int dfsSz(int i, int p, vector<vector<int> > & G, vector<int> & subTreeSz){
    int res=1;

    for (int j : G[i]){
        if (j == p) continue;

        res += dfsSz(j, i, G, subTreeSz);
    }

    return subTreeSz[i] = res;
}

void dfs(int i, int p, vector<vector<int> > & G, vector<vector<int> > & edgeId, vector<int> & subTreeSz, vector<vector<int> > & minstNode, vector<int> & minstSz, map<int,int> & minstNumNodes, int & diffMinst, vector<int> & res){
    int heavyJ=-1, heavyId;
    for (int ix=0; ix<(int)G[i].size(); ++ix){
        int j=G[i][ix], jId = edgeId[i][ix];

        if (j == p) continue;

        if (heavyJ == -1 || subTreeSz[j] > subTreeSz[heavyJ]){
            heavyJ = j;
            heavyId = jId;
        }
    }

    if (heavyJ == -1){
        diffMinst = 0;
        for (int m : minstNode[i]){
            minstNumNodes[m] = 1;
            diffMinst++;
        }
        
        return;
    }

    dfs(heavyJ, i, G, edgeId, subTreeSz, minstNode, minstSz, minstNumNodes, diffMinst, res);
    res[heavyId] = diffMinst;

    for (int ix=0; ix < (int)G[i].size(); ++ix){
        int j = G[i][ix], jId = edgeId[i][ix];

        if (j == heavyJ || j == p) continue;

        
        int subDiffMinst=0;
        map<int,int> subMinstNumNodes;
        dfs(j, i, G, edgeId, subTreeSz, minstNode, minstSz, subMinstNumNodes, subDiffMinst, res);
        
        res[jId] = subDiffMinst;

        for (auto [m, cnt] : subMinstNumNodes){
            if (cnt == -1) continue;

            if (minstNumNodes[m] == 0){
                minstNumNodes[m] = cnt;
                ++diffMinst;
            } else if (minstNumNodes[m] + cnt == minstSz[m]){
                minstNumNodes[m] = -1;
                --diffMinst;
            } else {
                minstNumNodes[m] += cnt;
            }
        }
    }

    for (int m : minstNode[i]){
        if (minstNumNodes[m] == 0){
            minstNumNodes[m] = 1;
            ++diffMinst;
        } else if (minstNumNodes[m] + 1 == minstSz[m]){
            minstNumNodes[m] = -1;
            --diffMinst;
        } else {
            minstNumNodes[m] += 1;
        }
    }
}

void solve(){
    int N, M, K; cin >> N >> M >> K;

    vector<vector<int> > G(N);
    vector<vector<int> > edgeId(N);
    for (int i=0; i<N-1; ++i){
        int u, v; cin >> u >> v;
        --u; --v;
        
        G[u].push_back(v);
        edgeId[u].push_back(i);

        G[v].push_back(u);
        edgeId[v].push_back(i);
    }

    vector<vector<int> > minstNode(N);
    vector<int> minstSz(M);
    for (int i=0; i<M; ++i){
        int s; cin >> s;
        minstSz[i] = s;

        for (int j=0; j<s; ++j){
            int u; cin >> u;
            --u;
            minstNode[u].push_back(i);
        }
    }

    int root = 0;
    vector<int> subTreeSz(N);
    dfsSz(root, -1, G, subTreeSz);

    vector<int> res(N-1);
    map<int,int> minstNumNodes;
    int diffMinst;
    dfs(root, -1, G, edgeId, subTreeSz, minstNode, minstSz, minstNumNodes, diffMinst, res);

    vector<int> ans;
    for (int i=0; i<N-1; ++i){
        if (res[i] >= K){
            ans.push_back(i+1);
        }
    }
    
    cout << ans.size() << '\n';
    for (int n : ans) cout << n << ' ';
    cout << '\n';
}

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

    solve();
    return 0;
}

• Subtask 1: 8 / 8
• Subtask 2: 15 / 15
• Subtask 3: 7 / 7
• Subtask 4: 29 / 29
• Subtask 5: 16 / 16
• Subtask 6: 25 / 25