Ice cream coloring | Codeforces #411 (Div. 1)

コメント付きコード

// 頂点の数、色の数、頂点vに含まれるアイスの数
int N, M, S[300001];
// グラフ, 頂点vに含まれるアイス
vi G[300001], C[300001];
// 解, 色iを使ったかどうか
int X[300001], use[300001];

int main(){
    scanf("%d%d", &N, &M);
    rep(i, N) {
        scanf("%d", &S[i]);
        C[i].resize(S[i]);
        rep(j, S[i])scanf("%d", &C[i][j]), --C[i][j];
    }
    rep(i, N - 1) {
        int u, v;
        scanf("%d%d", &u, &v);
        --u; --v;
        /*
        S[u] > 0 かつ S[v] > 0の場合だけ辺を張る。
        こうすると元の木Tの部分木の集合(Tの分割)が得られる。
        アイスは連結した部分グラフになっているので、どのアイスもたかだか1個の部分木にしか現れない。
        */
        if (S[u] && S[v]) {
            G[u].push_back(v);
            G[v].push_back(u);
        }
    }

    int ans = 1;
    queue<pii> q;
    rep(i, N) {
        /*
        頂点iの含む部分木に含まれるアイスの色をすべて決める。
        */
        if (S[i] && !X[C[i][0]]) {
            /*
            部分木を頂点ごとに探索していく。スタックオーバーフローを避けるためにDFSはしない。
            */
            q.push(mp(i, -1));
            while (sz(q)) {
                int u, pre;
                tie(u, pre) = q.front();
                q.pop();
                /*
                すでに色が確定しているアイスをチェック
                */
                rep(j, S[u]) {
                    int ice = C[u][j];
                    if (X[ice]) {
                        use[X[ice]] = 1;
                    }
                }
                /*
                1から順に、まだ使われていない色を探していく。
                */
                int cur = 1;
                rep(j, S[u]) {
                    int ice = C[u][j];
                    if (!X[ice]) {
                        // 色curは他のアイスに使われている。
                        while (use[cur])cur++;
                        X[ice] = cur;
                        use[cur] = 1;
                    }
                }
                /*
                curが最大値を取るとき、
                色1, 2, 3, 4, 5, ... , cur
                はすべて使われている。
                逆にこの時、curを一色でも減らすと、必ずどこかで同色が隣接してしまう。
                */
                smax(ans, cur);
                rep(j, S[u])use[X[C[u][j]]] = 0;
                each(v, G[u]) if (v != pre) {
                    q.push(mp(v, u));
                }
            }
        }
    }

    /*
    問題文には
    "empty set of vertices form a connected subgraph in this problem"
    とあるので、アイスが空の部分木になっている場合がある。
    このアイスの色は任意なので1をふっておく。
    人々が落ちてるのはここ？
    */
    rep(i, M)if (!X[i])X[i] = 1;

    printf("%d\n", ans);
    rep(i, M) {
        printf("%d%c", X[i], i != M - 1 ? ' ' : '\n');
    }
}