hogecoder

つたじろう(Tsuta_J) 競技プログラミングの記録

TopCoder SRM 716 Div1 Med (Div2 Hard): JumpDistancesOnTree

するめうめ

tsutaj.hatenablog.com

問題概要

原文 → TopCoder Statistics - Problem Statement

 N 頂点からなる木がある。頂点  0 を起点として、いくつかの頂点に移動することを考える。集合  S が与えられるので、頂点  v_{1} \rightarrow v_{2} \rightarrow \dots \rightarrow v_{k} の順に移動したときの隣り合う要素どうしの距離  d_{1}, d_{2}, \dots, d_{k-1} の中に  S の要素が全て含まれるような移動方法があるかを判定せよ。

解説

頂点  0 から行ける場所であって、頂点間の距離が  S の元であるような要素を管理したいです。これをするにはどうすればよいでしょうか?

 S に含まれている移動距離であるような頂点間が何であるかを把握します。これは全点対の距離を調べることで可能です。そして、その頂点間にのみ辺を貼ったグラフを新たに考え、 Union Find で連結成分を把握します。頂点  0 が属している連結成分に  S の元が全てあればよいので、それを判定すれば解くことが出来ます。

Div1 では制約が大きいだけなのですが、LCA 使った全点対を書いたら  O(N^{2} \log N) を通してくれないようで悲しい・・・もう少し考えます・・・

(追記) 木なら普通に dfs なりをやって全点対  O(N^{2}) だから LCA とかいらなかったですね (アホすぎる)

ソースコード

Div1 (UnionFind 省略, Div2 でもこれで通るはず)

class JumpDistancesOnTree {
    public:
    int dist[MAX_N + 10][MAX_N + 10];
    int assigned[MAX_N + 10];

    void dfs(Graph<int> &G, int cur, int orig, int par=-1, int d=0) {
        dist[cur][orig] = dist[orig][cur] = d;
        for(auto e : G[cur]) {
            if(e.to == par) continue;
            dfs(G, e.to, orig, cur, d+e.cost);
        }
    }

    string isPossible(vector<int> p, vector<int> S) {
        memset(assigned, false, sizeof(assigned));
        int N = p.size() + 1;
        for(int i=0; i<N; i++) {
            for(int j=0; j<N; j++) {
                dist[i][j] = (i == j ? 0 : INF);
            }
        }

        for(auto x : S) {
            assigned[x] = true;
        }

        Graph<int> G(N);
        for(int i=0; i<N-1; i++) {
            int u = i+1, v = p[i];
            G[u].push_back(Edge<int>(v, 1));
            G[v].push_back(Edge<int>(u, 1));
        }
        for(int i=0; i<N; i++) {
            dfs(G, i, i);
        }

        UnionFind uf;
        for(int i=0; i<N; i++) {
            for(int j=0; j<N; j++) {
                if(assigned[ dist[i][j] ]) {
                    uf.unite(i, j);
                }
            }
        }

        int ans = 0, root = uf.find(0);
        vector<int> cnt(MAX_N);
        for(int i=0; i<N; i++) {
            int idx = uf.find(i);
            if(idx != root) continue;
            for(int j=0; j<N; j++) {
                if(assigned[ dist[i][j] ]) {
                    if(!cnt[dist[i][j]]) ans++;
                    cnt[dist[i][j]] = true;
                }
            }
        }

        bool ok = (ans == (int)S.size());
        return (ok ? "Possible" : "Impossible");
    }
};