【洛谷】P3884二叉树问题

#include <iostream>
#include <cstdio>
#include <vector>
using namespace std;

//解题思路：
//倍增法打表求最近公共祖先
//倍增法原理：任何一个数都可以由二进制数组成、表示，则最近祖先的距离也可以分成2的幂次的各个段落

//其实这道题可以不用lca求解，因为数据较弱，只需求一组的lca
//所以也可以不用这种方法，反向建图dfs暴力求解应该也可以
inline void read(int &x)
{
    register int f = 1;
    x = 0;
    char c = getchar();
    while(c > '9' || c < '0')
    {
        if(c == '-') f = -1;
        c = getchar();
    }
    while(c >= '0' && c <= '9')
    {
        x = (x << 3) + (x << 1) + (c ^ 48);
        c  = getchar();
    }
    x *= f;
}

vector<int> mp[105];
// 深度和宽度数组，以及最大深度、宽度
int depth[105], wid[105], maxdep, maxwid;
// fa[i][j]数组表示2^j的祖先节点，即fa[i][0]则是i的父节点
int fa[105][8], lg[105], n;

// 求深度,宽度
void dfs(int u, int lev)
{
    wid[lev]++; depth[u] = lev;
    maxdep = max(lev, maxdep);
    if(!mp[u].size()) return;
    lev++;
    for(int i = 0; i < mp[u].size(); i++)
    {
        dfs(mp[u][i], lev);
    }
}

void init()
{
    dfs(1, 1);
//    预处理log2()的结果，顺便求出最大宽度
    for(int i = 1; i <= maxdep; i++)
    {
        lg[i] = lg[i - 1] + (1 << lg[i - 1] == i);
        maxwid = max(maxwid, wid[i]);
    }
//    打表
    for(int j = 1; j < 8; j++)
    {
        for(int i = 1; i <= n; i++)
        {
            fa[i][j] = fa[fa[i][j - 1]][j - 1];
        }
    }
}

//求最近公共祖先
int lca(int a, int b)
{
    if(depth[a] < depth[b]) swap(a,b);
    while(depth[a] > depth[b])
    {
        int len = depth[a] - depth[b];
        a = fa[a][lg[len] - 1];
    }
    if(a == b) return a;
    for(int i = lg[depth[a]] - 1; i >= 0; i--)
    {
        if(fa[a][i] != fa[b][i])
        {
            a = fa[a][i]; b = fa[b][i];
        }
    }
    return fa[a][0];
}

int main()
{
    read(n);
    for(int x, y, i = 1; i < n; i++)
    {
        read(x); read(y);
        mp[x].push_back(y);
        fa[y][0] = x;
    }
    init();

    printf("%d\n", maxdep);
    printf("%d\n", maxwid);

    int a, b, t, ans;
    read(a);read(b);

    t = lca(a,b);
    ans = (depth[a] - depth[t]) * 2 + (depth[b] - depth[t]);
    printf("%d\n", ans);
    return 0;
}