Codeforces 1182D Complete Mirror hash

Complete Mirror

嗯嗯嗯。。 调了一万年。

就是把树hash一下， 然后dfs两次求合法的根。

//#pragma GCC optimize(2)
//#pragma GCC optimize(3)
//#pragma GCC optimize(4)
#include<bits/stdc++.h>
#define LL long long
#define LD long double
#define ull unsigned long long
#define fi first
#define se second
#define mk make_pair
#define PLL pair<LL, LL>
#define PLI pair<LL, int>
#define PII pair<int, int>
#define SZ(x) ((int)x.size())
#define ALL(x) (x).begin(), (x).end()
#define fio ios::sync_with_stdio(false); cin.tie(0);

using namespace std;

const int N = 1e5 + 7;
const int inf = 0x3f3f3f3f;
const LL INF = 0x3f3f3f3f3f3f3f3f;
const int mod = 1e9 + 7;
const double eps = 1e-8;
const double PI = acos(-1);

template<class T, class S> inline void add(T &a, S b) {a += b; if(a >= mod) a -= mod;}
template<class T, class S> inline void sub(T &a, S b) {a -= b; if(a < 0) a += mod;}
template<class T, class S> inline bool chkmax(T &a, S b) {return a < b ? a = b, true : false;}
template<class T, class S> inline bool chkmin(T &a, S b) {return a > b ? a = b, true : false;}

mt19937 rng(chrono::steady_clock::now().time_since_epoch().count());

const LL B = 233333;
const LL A = 2333;
int n;
vector<int> G[N];

LL P;
LL hs[N];
LL hsson[N];
int ans;
bool can[N];

bool prime(LL x) {
    for(int i = 2; i * i <= x; i++) {
        if(x % i == 0) return false;
    }
    return true;
}

void dfs(int u, int fa) {
    can[u] = true;
    LL pre = -1;
    for(auto &v : G[u]) {
        if(v == fa) continue;
        dfs(v, u);
        can[u] &= can[v];
        if(pre == -1) pre = hs[v];
        else if(~pre && hs[v] != pre) {
            can[u] = false;
        }
        hs[u] = hs[u] * B % P;
        hs[u] = (hs[u] + hs[v]) % P;
    }
    hsson[u] = pre;
    hs[u] = hs[u] * B % P;
    hs[u] = (hs[u] + A) % P;
}

void dfs2(int u, int fa, LL hsfa) {
    if(~ans) return;
    if(can[u]) {
        if(hsson[u] == -1 || hsson[u] == hsfa) {
            ans = u;
            return;
        }
    }

    if(u != 1) {
        LL val1 = hsfa, val2 = -1;
        int cnt1 = 0, cnt2 = 0;
        int pos = -1;
        for(auto &v : G[u]) {
            if(v == fa) continue;
            if(!can[v]) {
                if(pos == -1) {
                    pos = v;
                } else {
                    return;
                }
                continue;
            }
            if(val1 == -1) {
                val1 = hs[v];
                cnt1++;
            }
            else if(hs[v] == val1) {
                cnt1++;
            }
            else if(val2 == -1) {
                val2 = hs[v];
                cnt2++;
            }
            else if(hs[v] == val2) {
                cnt2++;
            }
            else {
                return;
            }
        }
        if(~pos) {
            if(cnt2) return;
            LL nhsfa = hsfa;
            for(int i = 1; i <= cnt1; i++) {
                nhsfa = nhsfa * B % P;
                nhsfa = (nhsfa + val1) % P;
            }
            nhsfa = nhsfa * B % P;
            nhsfa = (nhsfa + A) % P;
            dfs2(pos, u, nhsfa);
            return;
        }
        if(cnt2 >= 2) return;
        if(cnt2) {
            LL nhsfa = hsfa;
            for(int i = 1; i <= cnt1; i++) {
                nhsfa = nhsfa * B % P;
                nhsfa = (nhsfa + val1) % P;
            }
            nhsfa = nhsfa * B % P;
            nhsfa = (nhsfa + A) % P;
            for(auto &v : G[u]) {
                if(v == fa || hs[v] != val2) continue;
                dfs2(v, u, nhsfa);
            }
        }
        else {
            LL nhsfa = hsfa;
            for(int i = 1; i < cnt1; i++) {
                nhsfa = nhsfa * B % P;
                nhsfa = (nhsfa + val1) % P;
            }
            nhsfa = nhsfa * B % P;
            nhsfa = (nhsfa + A) % P;
            for(auto &v : G[u]) {
                if(v == fa || hs[v] != val1) continue;
                dfs2(v, u, nhsfa);
            }
        }
    }
    else {
        LL val1 = -1, val2 = -1;
        int cnt1 = 0, cnt2 = 0;
        int pos = -1;
        for(auto &v : G[u]) {
            if(v == fa) continue;
            if(!can[v]) {
                if(pos == -1) {
                    pos = v;
                }
                else {
                    return;
                }
                continue;
            }
            if(val1 == -1) {
                val1 = hs[v];
                cnt1++;
            }
            else if(hs[v] == val1) {
                cnt1++;
            }
            else if(val2 == -1) {
                val2 = hs[v];
                cnt2++;
            }
            else if(hs[v] == val2) {
                cnt2++;
            }
            else {
                return;
            }
        }
        if(~pos) {
            if(cnt2) return;
            LL nhsfa = hsfa;
            for(int i = 1; i <= cnt1; i++) {
                nhsfa = nhsfa * B % P;
                nhsfa = (nhsfa + val1) % P;
            }
            nhsfa = nhsfa * B % P;
            nhsfa = (nhsfa + A) % P;
            dfs2(pos, u, nhsfa);
            return;
        }
        if(cnt1 && !cnt2) {
            LL nhsfa = hsfa;
            for(int i = 1; i < cnt1; i++) {
                nhsfa = nhsfa * B % P;
                nhsfa = (nhsfa + val1) % P;
            }
            nhsfa = nhsfa * B % P;
            nhsfa = (nhsfa + A) % P;
            for(auto &v : G[u]) {
                if(v == fa || hs[v] != val1) continue;
                dfs2(v, u, nhsfa);
            }
        }
        else {
            if(cnt1 > cnt2) {
                swap(cnt1, cnt2);
                swap(val1, val2);
            }
            if(cnt1 > 1) return;
            if(cnt2 > 1) {
                LL nhsfa = hsfa;
                for(int i = 1; i <= cnt2; i++) {
                    nhsfa = nhsfa * B % P;
                    nhsfa = (nhsfa + val2) % P;
                }
                nhsfa = nhsfa * B % P;
                nhsfa = (nhsfa + A) % P;
                for(auto &v : G[u]) {
                    if(v == fa || hs[v] != val1) continue;
                    dfs2(v, u, nhsfa);
                }
            }
            else {
                LL nhsfa1 = hsfa, nhsfa2 = hsfa;
                nhsfa1 = (nhsfa1 * B + val2) % P;
                nhsfa1 = (nhsfa1 * B + A) % P;
                nhsfa2 = (nhsfa2 * B + val1) % P;
                nhsfa2 = (nhsfa2 * B + A) % P;
                for(auto &v : G[u]) {
                    if(v == fa) continue;
                    if(hs[v] == val1) {
                        dfs2(v, u, nhsfa1);
                    }
                    else if(hs[v] == val2) {
                        dfs2(v, u, nhsfa2);
                    }
                }
            }
        }
    }
}

int main() {
    
    while(P <= 500000000) P = rng() % 900000000;
    while(!prime(P)) P++;

    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);
    }

    int Rt = 1;
    dfs(Rt, 0);
    
    if(can[Rt]) {
        printf("%d\n", Rt);
        return 0;
    }

    ans = -1;
    dfs2(Rt, 0, 0);

    printf("%d\n", ans);
    return 0;
}

/*
*/