POJ 1986 Distance Queries（查询两点距离，LCA）

题目链接：
POJ 1986 Distance Queries
题意：
给一个连通的$n$个节点的树，有$Q$次查询，每次输出两点间距离。
数据范围：$n\leq 40000$
分析：
先用基于RMQ算法的求LCA的方法求出LCA。记$dis[i]$为根节点到$i$节点的距离，那么$u$和$v$之间的距离就是：

$$Ans=dis[u]+dis[v] - 2 * dis[LCA(u,v)]$$
我建单向边wa了，双向边就AC了，不太懂啊。。。
时间复杂度：预处理:$O(n\log n)$，查询：$O(1)$

#include <stdio.h>
#include <string.h>
#include <algorithm>
#include <math.h>
using namespace std;
typedef long long ll;
const int MAX_N = 40010;

int n, m, total;
int head[MAX_N], in[MAX_N], id[MAX_N], dis[MAX_N];
int vis[MAX_N * 2], depth[MAX_N * 2], dp[MAX_N * 2][20];

struct Edge {
    int v, w, next;
}edge[MAX_N * 2];

void AddEdge (int u, int v, int w)
{
    edge[total].v = v, edge[total].w = w;
    edge[total].next = head[u];
    head[u] = total++;
}

void dfs(int u, int p, int d, int& k)
{
    vis[k] = u, id[u] = k;
    depth[k++] = d;
    for (int i = head[u]; i != -1; i = edge[i].next) {
        int v = edge[i].v, w = edge[i].w;
        if (v == p) continue;
        dis[v] = dis[u] + w;
        dfs(v, u, d + 1, k);
        vis[k] = u;
        depth[k++] = d;
    }
}

void RMQ(int root)
{
    int k = 0;
    dfs(root, -1, 0, k);
    int mm = k;
    int e = (int)log2(mm + 1.0);
    for (int i = 0; i < mm; ++i) dp[i][0] = i;
    for (int j = 1; j <= e; ++j) {
        for (int i = 0; i + (1 << j) - 1 < mm; ++i) {
            int nxt = i + (1 << (j - 1));
            if (depth[dp[i][j - 1]] < depth[dp[nxt][j - 1]]) {
                dp[i][j] = dp[i][j - 1]; 
            } else {
                dp[i][j] = dp[nxt][j - 1];
            }
        }
    }
}

int LCA(int u, int v)
{
    int left = min(id[u], id[v]), right = max(id[u], id[v]);
    int k = (int)log2(right - left + 1.0);
    int pos, nxt = right - (1 << k) + 1;
    if (depth[dp[left][k]] < depth[dp[nxt][k]]) {
        pos = dp[left][k];
    } else {
        pos = dp[nxt][k];
    }
    return dis[u] + dis[v] - 2 * dis[vis[pos]];
}

void init()
{
    total = 0;
    memset(head, -1, sizeof(head));
    memset(vis, 0, sizeof(head));
    memset(in, 0, sizeof(in));
    memset(dis, 0, sizeof(dis));
    for (int i = 0; i <= n; ++i) { fa[i] = i; }
}

int main()
{
    while (~scanf("%d%d", &n, &m)) {
        init();
        for (int i = 0; i < m; ++i) {
            int u, v, w;
            char s[10];
            scanf("%d%d%d%s", &u, &v, &w, s);
            AddEdge(u, v, w);
            AddEdge(v, u, w);
            in[v]++;
        }
        int root;
        for (int i = 1; i <= n; ++i) {
            if (in[i] == 0) {
                root = i;
                break;
            }
        }
        RMQ(root);
        scanf("%d", &m);
        while (m--) {
            int u, v;
            scanf("%d%d", &u, &v);
            printf("%d\n", LCA(u, v));
        }
    }
    return 0;
}