题解 | 小红的推理树平衡链路

//树上前缀和
//建树后手动模拟递归，用stage存储当前节点状态实现dfs
//presum表示从根节点到父节点的前缀和
//cursum表示       到当前节点的

#include <iostream>
#include <algorithm>
#include <vector>
#include <string>
#include <queue>
#include <unordered_map>
#include <stack>
using namespace std;
using ll = long long;

struct node
{
	int lchild;
	int rchild;
	ll w;
	bool isNone;
	node() :lchild(-1), rchild(-1), w(0), isNone(false) {}
};
vector<node> nodes;

void tree_init()
{
	int index = 1;
	queue<int> q;
	q.push(0);
	int n = (int)nodes.size();
	while (!q.empty() && index < n)
	{
		int cur = q.front();
		q.pop();
		if (nodes[cur].isNone) continue; //当前节点为空

		if (index < n && !nodes[index].isNone)
		{
			nodes[cur].lchild = index;
			q.push(index);
		}
		index++;
		if (index < n && !nodes[index].isNone)
		{
			nodes[cur].rchild = index;
			q.push(index);
		}
		index++;
	}
	return;
}

struct frame
{
	ll cursum;
	int stage; //0刚进入 1处理完左子树 2处理完右子树
	ll presum;
	int num;
};

unordered_map<ll, ll> mp;

ll func()
{
	ll ans = 0;
	stack<frame> st;

	st.push({ 0,0,0,0 });
	while (!st.empty())
	{
		frame& tmp = st.top();
		//cout << "tmp.num:" << tmp.num << endl;
		//st.pop();
		
		if (nodes[tmp.num].isNone)
		{
			st.pop();
			continue;
		}
		
		tmp.cursum = tmp.presum + nodes[tmp.num].w; 
		if (tmp.stage == 0)
		{
			ans += mp[tmp.cursum]; //ans+(count(mp[祖先节点前缀和]==mp[当前节点前缀和]))
			mp[tmp.presum]++; //该前缀和的祖先节点个数+1
			tmp.stage = 1;
			//st.push(tmp);
			if (nodes[tmp.num].lchild != -1)
			{
				st.push({0, 0, tmp.cursum, nodes[tmp.num].lchild});
			}
		}
		else if (tmp.stage == 1)
		{
			tmp.stage = 2;
			//st.push(tmp);
			if (nodes[tmp.num].rchild != -1)
			{
				st.push({ 0, 0, tmp.cursum, nodes[tmp.num].rchild });
			}
		}
		else
		{
			mp[tmp.presum]--;
			st.pop();
		}
	}
	return ans;
}

int main()
{
	ios::sync_with_stdio(0);
	cin.tie(0), cout.tie(0);
	int n;
	cin >> n;
	nodes.resize(n);
	for (int i = 0; i < n; i++)
	{
		string s;
		cin >> s;
		if (s != "None")
		{
			nodes[i].w = stoll(s);
		}
		else
		{
			nodes[i].isNone = true;
		}
	}
	tree_init();
	ll ans = func();
	cout << ans << endl;
	return 0;
}