Codeforces 1073F Choosing Two Paths 树形dp

Choosing Two Paths

搞不懂为啥我要写换根啊， 这么麻烦。。

#include<bits/stdc++.h>
#define fi first
#define se second
#define mk make_pair
#define PII pair<int, int>
using namespace std;

const int N = (int)2e5 + 7;

int n, pa[N];
vector<int> G[N];

struct Node {
    PII one;
    pair<PII, PII> two;
    inline Node up() {
        Node ret = *this;
        if(ret.one.se) ret.one.fi++;
        if(ret.two.se.fi) ret.two.fi.fi++, ret.two.fi.se += 2;
        return ret;
    }
    void show() {
        printf("(%d, %d)  ((%d, %d) (%d, %d))\n", one.fi, one.se, two.fi.fi, two.fi.se, two.se.fi, two.se.se);
    }
} dp[N], fdp[N];

inline void update(Node &A, Node B) {
    if(B.two.fi > A.two.fi) A.two = B.two;
    if(A.one.se && B.one.se && mk(1, A.one.fi + B.one.fi) > A.two.fi) A.two = mk(mk(1, A.one.fi + B.one.fi), mk(A.one.se, B.one.se));
    if(B.one.se && B.one.fi > A.one.fi) A.one = B.one;
}

void dfs(int u, int fa) {
    pa[u] = fa;
    if(u != 1 && (int)G[u].size() == 1) {
        dp[u].one = mk(0, u);
        return;
    }
    for(auto &v : G[u]) {
        if(v == fa) continue;
        dfs(v, u);
        update(dp[u], dp[v].up());
    }
}

void dfs2(int u, int fa) {
    if(u != 1 && G[u].size() == 1) {
        return;
    }
    vector<Node> pre(G[u].size());
    vector<Node> suf(G[u].size());
    for(int i = 0; i < G[u].size(); i++) {
        int v = G[u][i];
        if(v == fa) {
            if(i) pre[i] = pre[i - 1], update(pre[i], fdp[u].up());
            else pre[i] = fdp[u].up();
        }
        else {
            if(i) pre[i] = pre[i - 1], update(pre[i], dp[v].up());
            else pre[i] = dp[v].up();
        }
    }
    for(int i = (int)G[u].size() - 1; i >= 0; i--) {
        int v = G[u][i];
        if(v == fa) {
            if(i != (int)G[u].size() - 1) suf[i] = suf[i + 1], update(suf[i], fdp[u].up());
            else suf[i] = fdp[u].up();
        }
        else {
            if(i != (int)G[u].size() - 1) suf[i] = suf[i + 1], update(suf[i], dp[v].up());
            else suf[i] = dp[v].up();
        }
    }
    if((int)G[u].size() == 1) {
        fdp[G[u][0]].one = mk(0, u);
        dfs2(G[u][0], u);
        return;
    }
    for(int i = 0; i < G[u].size(); i++) {
        int v = G[u][i];
        if(v == fa) continue;
        if(i) {
            fdp[v] = pre[i - 1];
            if(i < (int)G[u].size() - 1) update(fdp[v], suf[i + 1]);
        }
        else {
            fdp[v] = suf[i + 1];
        }
        dfs2(v, u);
    }
}

int main() {
    scanf("%d", &n);
    for(int i = 1; i < n; i++) {
        int u, v;
        scanf("%d%d", &u, &v);
        G[u].push_back(v);
        G[v].push_back(u);
    }
    dfs(1, 0);
    dfs2(1, 0);
    PII ret = mk(0, 0);
    PII p1, p2;
    for(int u = 1; u <= n; u++) {
        vector<Node> V;
        for(auto &v : G[u]) {
            if(v == pa[u]) V.push_back(fdp[u]);
            else V.push_back(dp[v]);
        }
        pair<PII, int> mx0 = mk(mk(-1, -1), -1);
        pair<PII, int> mx1 = mk(mk(-1, -1), -1);
        PII mx[3];
        for(int i = 0; i < 3; i++) mx[i] = mk(-1, -1);
        for(int i = 0; i < (int)V.size(); i++) {
            if(V[i].two.se.fi) {
                if(V[i].two.fi > mx0.fi) mx1 = mx0, mx0 = mk(V[i].two.fi, i);
                else if(V[i].two.fi > mx1.fi) mx1 = mk(V[i].two.fi, i);
            }
            if(V[i].one.se) {
                PII cur = mk(V[i].one.fi, i);
                for(int j = 0; j < 3; j++) {
                    if(cur > mx[j]) swap(cur, mx[j]);
                }
            }
        }
        if(mx1.se != -1) {
            PII tmp = mk(mx0.fi.fi + mx1.fi.fi + 1, mx0.fi.se + mx1.fi.se + 4);
            if(tmp > ret) {
                ret = tmp;
                p1 = V[mx0.se].two.se;
                p2 = V[mx1.se].two.se;
            }
        }
        for(int i = 0; i < V.size(); i++) {
            if(!V[i].two.se.fi) continue;
            PII mx0 = mk(-1, -1);
            PII mx1 = mk(-1, -1);
            for(int j = 0; j < 3; j++) {
                if(mx[j].se == -1 || mx1.se != -1) break;
                if(mx[j].se == i) continue;
                if(mx0.se == -1) mx0 = mx[j];
                else mx1 = mx[j];
            }
            if(mx1.se != -1) {
                PII tmp = mk(V[i].two.fi.fi + 1, V[i].two.fi.se + mx0.fi + mx1.fi + 4);
                if(tmp > ret) {
                    ret = tmp;
                    p1 = V[i].two.se;
                    p2 = mk(V[mx0.se].one.se, V[mx1.se].one.se);
                }
            }
        }
    }
    printf("%d %d\n%d %d\n", p1.fi, p2.fi, p1.se, p2.se);
    return 0;
}