USACO 2015 December Contest, Platinum Problem 1. Max Flow

http://www.usaco.org/index.php?page=viewproblem2&cpid=576

key : LCA

#include <iostream>
#include <stdio.h>
#include <vector>
#include <utility>
using namespace std;

typedef pair<int,int> pii;
const int MAX_N = 5e4+6;
const int MAX_K = 1e5+6;
const int MAX_P = 21;

pii pip[MAX_K]; //pipe
int ans[MAX_N];
int depth[MAX_N];
int lca[MAX_P][MAX_N];
int _1[MAX_N],_2[MAX_N];
int st[MAX_N];
int binary[MAX_P];
bool visit[MAX_N];
vector<int> edg[MAX_N];

void dfs1(int id,int cur_depth) {
    depth[id]=cur_depth;
    for (auto i=edg[id].begin();i!=edg[id].end();i++) {
        int tmp=*i;
        if (depth[tmp]==0) {
            lca[1][tmp]=id;
            dfs1(tmp,cur_depth+1);
        }
    }
}

int walk(int i,int to_depth) {
    int tmp=depth[i]-to_depth;
    if (tmp<0) return -1;
    int id=0;
    while (tmp>0) {
        binary[id++]=tmp%2;
        tmp/=2;
    }
    for (int x=id-1;x>=0;x--) {
        if (binary[x]==1) {
            i=lca[x+1][i];
        }
    }
    return i;
}

void get_lca(pii i) {
    int a=i.first;
    int b=i.second;
    st[a]++,st[b]++;
    int R=min(depth[a],depth[b])+1,L=0; //L is the answer
    while (R-L != 1) {
//        cout<<"R = "<<R<<" , L="<<L<<endl;
        int mid=(L+R)>>1;
        if (walk(a,mid) == walk(b,mid) && walk(a,mid)!=-1) L=mid;
        else R=mid;
    }
    int step=L;
//    cout<<"a = "<<a<<" , b="<<b<<" , step = "<<step<<" , to = "<<walk(a,step)<<endl;
    if (step==depth[a]) {
//        cout<<"in1\n";
        st[a]--;
        _2[a]++;
    }
    else if (step==depth[b]) {
//        cout<<"in2\n";
        st[b]--;
        _2[b]++;
    }
    else {
        _1[walk(a,step)]++;
    }
//    cout<<"a="<<a<<" , b="<<b<<" , meet = "<<walk(a,step)<<endl;
}

int dfs2(int id) {
    visit[id]=true;
    int ret=0;
    for (auto i=edg[id].begin();i!=edg[id].end();i++) {
        int tmp=*i;
        if (visit[tmp]==false) {
            ret += dfs2(tmp);
        }
    }
//    cout<<"id = "<<id<<" , st="<<st[id]<<" , _1="<<_1[id]<<" , _2="<<_2[id];
//    cout<<" ret = "<<ret;
    ans[id]=ret - _1[id] + st[id];
    ret = ret - _2[id] + st[id] - 2*_1[id];
//    cout<<" ans = "<<ans[id]<<endl;
    return ret;
}

int main () {
    if (fopen("maxflow.in","r")) {
        freopen("maxflow.in","r",stdin);
        freopen("maxflow.out","w",stdout);
    }
    int n,k;
    while (scanf("%d %d",&n,&k) != EOF) {
        for (int x=0;n>=x;x++) {
            edg[x].clear();
            depth[x]=st[x]=_1[x]=_2[x]=0;
            visit[x]=false;
        }
        for (int x=0;n-1>x;x++) {
            int i,j;
            scanf("%d %d",&i,&j);
            edg[j].push_back(i);
            edg[i].push_back(j);
        }
        for (int x=0;k>x;x++) {
            int i,j;
            scanf("%d %d",&i,&j);
            pip[x]=make_pair(i,j);
        }
        dfs1(1,1);
        lca[1][1]=1;
        for (int x=2;MAX_P>x;x++) {
            for (int y=1;n>=y;y++) {
                lca[x][y] = lca[x-1][lca[x-1][y]];
            }
        }
        for (int x=0;k>x;x++) {
            get_lca(pip[x]);
        }
        dfs2(1);
        int mx=-1;
        for (int x=1;n>=x;x++) {
//            cout<<"ans = "<<ans[x]<<endl;
            mx = max(mx,ans[x]);
        }
        printf("%d\n",mx);
    }
}