AOJ 0575 Festivals in JOI Kingdom

問題

http://judge.u-aizu.ac.jp/onlinejudge/description.jsp?id=0575

解法

実装できなかったので，eagletmt さんの解法を参考にした．
AOJ 0575 - Festivals in JOI Kingdom - プログラミングコンテストの記録

使うアルゴリズムとしては，DijkstraとKruskal，LCAの3つ．
まず，各街がお祭りをする街からどれくらい離れているかは，Dijkstra法で求められる．
街 i がお祭りしている街から離れている距離を d[i] とする．
次に，与えられたグラフ上の辺 e = (u, v) のコストを min(d[u], d[v]) と改める．
こうしてできたグラフ上で最大全域木問題（つまり辺のコストが大きい方から追加していく）を解く．これで木ができる．
すると，任意の2つの街 u, v の組に対し，u から v への移動の際には今構築した最大全域木上を動けば良いことがわかる．
あとは u から v へのパスのなかで一番値が小さいものを効率よく求められれば良い．

ここで，最大全域木を構築するときに工夫をすると，LCAでそれが楽に求められる．
Union-find木で街 u と街 v を合併していくわけだが，この過程を木として残しておく．
これは，合併の度に新しいノード new_v を作ることで実現する．街 u と街 v の合併のとき，ru = uf.root(u) と rv = uf.root(v) を，その新しいノード new_v にしてしまう．この際，new_v には，min(d[ru], d[rv]) を持たせておく．
そして，new_v から ru, rv に対し有向辺を張れば，これはまた木になっている．
各クエリは，こうしてできた木にたいして lca を求め，その値を出力すれば良い．

計算量は（ネックになるやつだけ見ると） O(ElogV + ElogE + QlogN) ．

手元で試すとスタックオーバーフローした．AOJだといけた．割りとギリギリ．
自分の書き方だとメモリ効率が悪くなってるかもしれない．

ソースコード

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

using P = pair<int, int>;

// コードが長いので，union_find は省略

struct edge {
    int from, to, cost;
    bool operator<(edge const& other) const {
        return cost > other.cost;
    }
};

using edges = vector<edge>;
using graph = vector<edges>;


class lca {
public:
    lca(vector<vector<int>>& g, int root)
        : V(g.size()),
          log_V(0),
          depth_(V, 0)
    {
        for(int v=V; v>0; v /= 2) {
            ++log_V;
        }
        parent.assign(log_V, std::vector<int>(V, 0));
        dfs(g, root, -1, 0);
        for(int k=0; k+1 < log_V; ++k) {
            for(int v=0; v<V; ++v) {
                if(parent[k][v] < 0) {
                    parent[k+1][v] = -1;
                } else {
                    parent[k+1][v] = parent[k][parent[k][v]];
                }
            }
        }
    }

    void dfs(vector<vector<int>> const& g, int v, int p, int d) {
        parent[0][v] = p;
        depth_[v] = d;
        for(int i=0; i<g[v].size(); ++i) {
            if(g[v][i] != p) {
                dfs(g, g[v][i], v, d+1);
            }
        }
    }

    int depth(int v) const {
        return depth_[v];
    }

    int query(int u, int v) {
        if(depth_[u] > depth_[v]) {
            std::swap(u, v);
        }
        for(int k=0; k<log_V; ++k) {
            if((depth_[v] - depth_[u]) >> k & 1) {
                v = parent[k][v];
            }
        }
        if(u == v) {
            return u;
        }
        for(int k=log_V - 1; k>=0; --k) {
            if(parent[k][u] != parent[k][v]) {
                u = parent[k][u];
                v = parent[k][v];
            }
        }
        return parent[0][u];
    }

private:
    int V, log_V;
    std::vector<vector<int>> parent;
    std::vector<int> depth_;
};


constexpr int INF = 1e9;

int main() {
    int N, M, K, Q;
    cin >> N >> M >> K >> Q;
    graph g(N);
    for(int i=0; i<M; ++i) {
        int a, b, l;
        cin >> a >> b >> l;
        a--; b--;
        g[a].push_back((edge){a, b, l});
        g[b].push_back((edge){b, a, l});
    }

    vector<int> d(2*N, INF);
    priority_queue<P, vector<P>, greater<P>> que;
    for(int i=0; i<K; ++i) {
        int f;
        cin >> f;
        f--;
        d[f] = 0;
        que.push(make_pair(0, f));
    }
    while(!que.empty()) {
        P p = que.top();
        que.pop();
        int v = p.second;
        if(d[v] < p.first) {
            continue;
        }
        for(auto& e : g[v]) {
            if(d[e.to] > d[v] + e.cost) {
                d[e.to] = d[v] + e.cost;
                que.push(make_pair(d[e.to], e.to));
            }
        }
    }

    edges es;
    for(int i=0; i<N; ++i) {
        for(auto& e : g[i]) {
            if(i < e.to) {
                es.push_back((edge){i, e.to, min(d[i], d[e.to])});
            }
        }
    }
    sort(es.begin(), es.end());

    union_find uf(N);
    vector<vector<int>> tree(2*N);
    int new_id = N;
    vector<int> id(N);
    for(int i=0; i<N; ++i) {
        id[i] = i;
    }
    for(auto& e : es) {
        int a = uf.root(e.from), b = uf.root(e.to);
        if(uf.unite(a, b)) {
            tree[new_id].push_back(id[a]);
            tree[new_id].push_back(id[b]);
            d[new_id] = min(d[id[a]], d[id[b]]);
            id[uf.root(a)] = new_id;
            new_id += 1;
        }
    }

    lca l(tree, new_id-1);
    for(int i=0; i<Q; ++i) {
        int s, t;
        cin >> s >> t;
        s--; t--;
        cout << d[l.query(s, t)] << endl;
    }
}