【算法】P5018 对称二叉树

题目

P5018 对称二叉树
https://www.luogu.com.cn/problem/P5018

代码

思路：领接表存储二叉树，unordered_map存储各个节点对应的值。dfs遍历一下各个子树的大小个数，再写个递归判断是否是对称二叉树，如果是就更新全局答案。

#include <bits/stdc++.h>
#define endl '\n'

using namespace std;

const int N = 1e7 + 10;

// 对于每个点k，开一个单链表，存储k所有可以走到的点。h[k]存储这个单链表的头结点（也就是邻接表）
int e[N], ne[N], h[N], st1[N], idx;

unordered_map<int, int> mp;// 每个结点id对应的值
int max_n = 0; // 最大对称二叉树节点数量

// 邻接表初始化
void init() {
	memset(h, -1, sizeof h);
	idx = 0;
}

// 添加一条边a->b 
void add(int a, int b) {
    // 存下b的值，b下一个指向a的下一个节点，a的下一个节点指向b
	e[idx] = b, ne[idx] = h[a], h[a] = idx ++;
}

//p, q是指针
bool check(int p, int q) {
	if (mp[e[p]] == 0 && mp[e[q]] == 0) // 递归到结尾返回true
		return true;
	else if (mp[e[p]] == 0 || mp[e[q]] == 0) // 两个孩子有一个为空返回false
		return false;
	else if (mp[e[p]] != mp[e[q]]) // 左孩子和右孩子不相同返回false
		return false;
	return check(h[e[p]], ne[h[e[q]]]) && check(ne[h[e[p]]], h[e[q]]); // 左右两颗子树开始递归
}

int dfs(int u) {
	if (mp[u] == 0) // 没有节点，返回0
		return 0;
	st1[u] = true;// st[u] 表示点u已经被遍历过
	int sum = 0;
	for (int i = h[u]; i != -1 ; i = ne[i]) {
		int j = e[i];
		if (st1[j]) continue;// 防止逆向dfs
		int s = dfs(j);
		sum += s; // 累加左孩子右孩子节点数
	}
	// 检查是不是对称二叉树，并更新答案
	if (check(h[u], ne[h[u]])) {
		max_n = max(max_n, sum + 1);
	}
	return sum + 1; // 返回当前节点的左孩子右孩子所有结点数+1
}

int main() {
	cin.tie(0), cout.tie(0);
	init();
	mp[-1] = 0;
	int n;
	cin >> n;
	// 每个节点存下节点值
	for (int i = 1; i <= n; i ++) {
		int v;
		cin >> v;
		mp[i] = v;
	}
	// 插入左孩子右孩子
	for (int i = 1; i <= n; i ++) {
		int l, r;
		cin >> l >> r;
		add(i, r);
		add(i, l);
	}
	// 从1开始dfs
	dfs(1);
	cout << max_n << endl;
	return 0;
}