AOJ 0575: Festivals in JOI Kingdom

Dijkstra(O(M log N)して、Maximum Spanning Treeを作ってO(M log N), Q個のクエリにLCAを使って応答をします(O(QlogN)).

全体のオーダーはO((M+Q)logN)くらいです. 詳しい解説はiwi先生のスライドなどを参考にしてください (多分想定解法そのままなので).

#include <cstdio>
#include <queue>
#include <vector>
#include <algorithm>

#define INF (1000000000)
#define MAX_V (100000)

using namespace std;

typedef struct {
    int to;
    int cost;
} EDGE;

typedef struct {
    int p, q, r;
} LOG;

typedef pair<int, int> P;

vector<EDGE> G[MAX_V];
vector<LOG> all;

/* Dijkstra */

int d[MAX_V];
priority_queue<P, vector<P>, greater<P> > que;

void dijkstra()
{
    
    while (!que.empty()){
        P p = que.top(); que.pop();
        int v = p.second;
        
        for (int i = 0; i < G[v].size(); i++){
            EDGE e = G[v][i];
            if (d[e.to] > d[v] + e.cost){
                d[e.to] = d[v] + e.cost;
                que.push(P(d[e.to], e.to));
            }
        }
    }
}

/* UnionFind */

int par[MAX_V];
int rank[MAX_V];

void init(int n)
{
    int i;
    for (i = 0; i < n; i++){
        par[i] = i;
        rank[i] = 0;
    }
}

int find(int x)
{
    if (par[x] == x){
        return (x);
    }
    return (par[x] = find(par[x]));
}

int same(int x, int y)
{
    return (find(y) == find(x));
}

void merge(int x, int y)
{
    x = find(x);
    y = find(y);
    
    if (x == y){
        return;
    }
    
    if (rank[x] < rank[y]){
        par[x] = y;
    }
    else {
        par[y] = x;
        if (x == y){
            rank[x]++;
        }
    }
}

bool comp(const LOG &a, const LOG &b)
{
    return (a.r > b.r);
}

/* LCA */

int root;

EDGE parent[20][MAX_V];
int depth[MAX_V];

void dfs(int v, int p, int dpth)
{
    parent[0][v].to = p;
    parent[0][v].cost = min(d[v], (p != -1 ? d[p] : INF));
    depth[v] = dpth;
    for (int i = 0; i < G[v].size(); i++){
        if (G[v][i].to != p){
            dfs(G[v][i].to, v, dpth + 1);
        }
    }
}

void initTree(int V)
{
    dfs(root, -1, 0);
    
    for (int k = 0; k + 1 < 20; k++){
        for (int v = 0; v < V; v++){
            if (parent[k][v].to < 0){
                parent[k + 1][v].to = -1;
            }
            else {
                parent[k + 1][v].to = parent[k][parent[k][v].to].to;
                parent[k + 1][v].cost = min(parent[k][parent[k][v].to].cost, parent[k][v].cost);
            }
        }
    }
}

int query(int u, int v)
{
    int res;
    res = min(d[u], d[v]);
    
    if (depth[u] > depth[v]){
        swap(u, v);
    }
    for (int k = 0; k < 20; k++){
        if (((depth[v] - depth[u]) >> k) & 1){
            res = min(res, parent[k][v].cost);
            v = parent[k][v].to;
        }
    }
    
    if (u == v){
        return (min(res, d[u]));
    }
    
    for (int k = 19; k >= 0; k--){
        if (parent[k][u].to != parent[k][v].to){
            res = min(res, min(parent[k][u].cost, parent[k][v].cost));
            u = parent[k][u].to;
            v = parent[k][v].to;
        }
    }
    
    return (min(res,  parent[0][v].cost));
}

int main(void)
{
    int N, M, K, Q;
    int from, go, m;
    int ans;
    
    scanf("%d%d%d%d", &N, &M, &K, &Q);
    
    fill(d, d + N, INF);
    
    for (int i = 0; i < M; i++){
        EDGE add;
        LOG mem;
        scanf("%d%d%d", &from, &go, &m);
        add.cost = m;
        add.to = --go;
        G[--from].push_back(add);
        add.to = from;
        G[go].push_back(add);
        mem.p = from;
        mem.q = go;
        mem.r = m;
        all.push_back(mem);
    }
    
    for (int i = 0; i < K; i++){
        scanf("%d", &from);
        d[--from] = 0;
        que.push(P(0, from));
    }
    dijkstra();
    
    for (int i = 0; i < N; i++){
        G[i].clear();
    }
    
    vector<LOG>::iterator it;
    
    for (it = all.begin(); it != all.end(); it++){
        it->r = min(d[it->p], d[it->q]);
    }
    
    sort(all.begin(), all.end(), comp);
    
    init(N);
    for (it = all.begin(); it != all.end(); it++){
        if (!same(it->p, it->q)){
            EDGE add;
            merge(it->p, it->q);
            add.to = it->q;
            add.cost = it->r;
            G[it->p].push_back(add);
            add.to = it->p;
            G[it->q].push_back(add);
        }
    }
    
    initTree(N);
    
    for (int i = 0; i < Q; i++){
        scanf("%d%d", &from, &go);
        printf("%d\n", query(--from, --go));
    }
    
    return (0);
}