python/TangDou/CSP-J/AcWingTraining/T9/476.cpp

#include <bits/stdc++.h>
using namespace std;

const int N = 1000010;

int n, w[N], l[N], r[N], si[N];

int getSize(int u) {
    if (!u) return 0;
    if (si[u]) return si[u];
    si[u] = 1 + getSize(l[u]) + getSize(r[u]);
    return si[u];
}

bool check(int u, int v) {
    if (u == 0 && v == 0)
        return true; // 两个空结点
    else if (w[u] != w[v])
        return false; // 权值不相等
    else if (si[u] != si[v])
        return false; // 树的大小不同
    // u树左子树和v右子树相同 并且u树右子树和v左子树相同
    else
        return check(l[u], r[v]) && check(r[u], l[v]);
}

int dfs(int u) {
    if (!u) return 0; // u树为空
    int ans = 0;

    // u树的左右子树对称
    if (check(l[u], r[u]))
        ans = si[u];
    else
        ans = max(dfs(l[u]), dfs(r[u]));

    return ans;
}

int main() {
    cin >> n;

    for (int i = 1; i <= n; i++) cin >> w[i];

    for (int i = 1; i <= n; i++) {
        int u, v;
        cin >> u >> v;
        // 为了方便处理数组下表，如果左右儿子为空，则为0
        if (~u) l[i] = u;
        if (~v) r[i] = v;
    }

    // 递归算出树中以每个节点为根的子树大小
    getSize(1);

    cout << dfs(1) << endl;

    return 0;
}