LeetCode834 树中距离之和

给定一个无向、连通的树。树中有 n 个标记为 0…n-1 的节点以及 n-1 条边 。
给定整数 n 和数组 edges ， edges[i] = [ai, bi]表示树中的节点 ai 和 bi 之间有一条边。
返回长度为 n 的数组 answer ，其中 answer[i] 是树中第 i 个节点与所有其他节点之间的距离之和。
示例 1:
输入: n = 6, edges = [[0,1],[0,2],[2,3],[2,4],[2,5]]
输出: [8,12,6,10,10,10]
解释: 树如图所示。
我们可以计算出 dist(0,1) + dist(0,2) + dist(0,3) + dist(0,4) + dist(0,5)
也就是 1 + 1 + 2 + 2 + 2 = 8。 因此，answer[0] = 8，以此类推。
示例 2:
输入: n = 1, edges = []
输出: [0]
示例 3:
输入: n = 2, edges = [[1,0]]
输出: [1,1]
参数:
1 <= n <= 3 * 10⁴
edges.length == n - 1
edges[i].length == 2
0 <= ai, bi < n
ai != bi
给定的输入保证为有效的树

深度优先搜索

假定节点0是树的根。
一，通过深度优先搜索计算m_vSubNodeCounts[i] ，以i为根节点的子树的节点数量。
二，通过深度优先搜索计算0到所有节点的距离。
DFSDis返回值是： cur到所有子孙节点的距离。
三，深度优先搜索，通过父节点到各点距离，计算当前节点到各点距离。
当前节点 和 当前节点的子孙节点 到 当前节点 的距离比到当前节点的父节点少1。
其它节点 到当前节点的距离比到当前节点的父节点的距离多1。

代码

核心代码

class CNeiBo2
{
public:
	CNeiBo2(int n, bool bDirect, int iBase = 0) :m_iN(n), m_bDirect(bDirect), m_iBase(iBase)
	{
		m_vNeiB.resize(n);
	}
	CNeiBo2(int n, vector<vector<int>>& edges, bool bDirect, int iBase = 0) :m_iN(n), m_bDirect(bDirect), m_iBase(iBase)
	{
		m_vNeiB.resize(n);
		for (const auto& v : edges)
		{
			m_vNeiB[v[0] - iBase].emplace_back(v[1] - iBase);
			if (!bDirect)
			{
				m_vNeiB[v[1] - iBase].emplace_back(v[0] - iBase);
			}
		}
	}
	inline void Add(int iNode1, int iNode2)
	{
		iNode1 -= m_iBase;
		iNode2 -= m_iBase;
		m_vNeiB[iNode1].emplace_back(iNode2);
		if (!m_bDirect)
		{
			m_vNeiB[iNode2].emplace_back(iNode1);
		}
	}
	const int m_iN;
	const bool m_bDirect;
	const int m_iBase;
	vector<vector<int>> m_vNeiB;
};


class Solution {
public:
	vector<int> sumOfDistancesInTree(int n, vector<vector<int>>& edges) {
		m_vSubNodeCounts.resize(n);
		m_vRet.resize(n);
		CNeiBo2 neiBo(n, edges, false);
		DFSSubNodeCount(neiBo, 0, -1);
		m_vRet[0] = DFSDis(neiBo, 0, -1);
		DFS(neiBo, 0, -1);
		return m_vRet;
	}
	void DFS(const CNeiBo2& neiBo, int cur, int parent)
	{
		if (-1 != parent)
		{
			m_vRet[cur] = m_vRet[parent] - m_vSubNodeCounts[cur] + (m_vSubNodeCounts.size() - m_vSubNodeCounts[cur]);
		}
		for (const auto& next : neiBo.m_vNeiB[cur])
		{
			if (parent == next)
			{
				continue;
			}
			DFS(neiBo, next, cur);
		}
	}
	int DFSSubNodeCount(const CNeiBo2& neiBo,int cur,int parent)
	{
		int iRet = 1;
		for (const auto& next : neiBo.m_vNeiB[cur])
		{
			if (parent == next)
			{
				continue;
			}
			iRet += DFSSubNodeCount(neiBo, next, cur);
		}
		return m_vSubNodeCounts[cur] = iRet;
	}
	int DFSDis(const CNeiBo2& neiBo, int cur, int parent)	
	{
		int iDis = m_vSubNodeCounts[cur]-1;
		for (const auto& next : neiBo.m_vNeiB[cur])
		{
			if (parent == next)
			{
				continue;
			}
			iDis += DFSDis(neiBo, next, cur);
		}
		return iDis;
	}
	vector<int> m_vSubNodeCounts,m_vRet;
};

测试用例

template<class T>
void Assert(const T& t1, const T& t2)
{
	assert(t1 == t2);
}

template<class T>
void Assert(const vector<T>& v1, const vector<T>& v2)
{
	if (v1.size() != v2.size())
	{
		assert(false);
		return;
	}
	for (int i = 0; i < v1.size(); i++)
	{
		Assert(v1[i], v2[i]);
	}

}

int main()
{
	int n;
	vector<vector<int>> edges;
	{
		Solution sln;
		n = 6, edges = { {0,1},{0,2},{2,3},{2,4},{2,5} };
		auto res = sln.sumOfDistancesInTree(n, edges);
		Assert(vector<int>{8, 12, 6, 10, 10, 10}, res);
	}
	{
		Solution sln;
		n = 1, edges = {};
		auto res = sln.sumOfDistancesInTree(n, edges);
		Assert(vector<int>{0}, res);
	}
	{
		Solution sln;
		n = 2, edges = { {1,0} };
		auto res = sln.sumOfDistancesInTree(n, edges);
		Assert(vector<int>{1,1}, res);
	}
	{
		Solution sln;
		n = 3, edges = { {2,1},{0,2} };
		auto res = sln.sumOfDistancesInTree(n, edges);
		Assert(vector<int>{3,3,2}, res);
	}
}

2023年1月

class Solution {
public:
vector sumOfDistancesInTree(int n, vector<vector>& edges) {
m_n = n;
m_vChildDis.assign(n, -1);
m_vChildNum.resize(n, -1);
m_vNP.resize(n);
m_vTotalNum.resize(n);
for (auto& e : edges)
{
m_vNP[e[0]].push_back(e[1]);
m_vNP[e[1]].push_back(e[0]);
}
dfs1(0, -1);
dfs2(0, -1);
return m_vTotalNum;
}
void dfs1(int iCur, const int iParent)
{
int iDis = 0;
int iNum = 0;
for (const auto& next : m_vNP[iCur])
{
if (iParent == next)
{
continue;
}
dfs1(next, iCur);
iDis += m_vChildDis[next];
iNum += m_vChildNum[next];
}
m_vChildDis[iCur] = iDis + iNum;
m_vChildNum[iCur] = iNum+1;
}
void dfs2(int iCur, const int iParent)
{
if (-1 == iParent)
{
m_vTotalNum[iCur] = m_vChildDis[iCur];
}
else
{
m_vTotalNum[iCur] = m_vTotalNum[iParent] - m_vChildDis[iCur] - m_vChildNum[iCur] + (m_n - m_vChildNum[iCur]) + m_vChildDis[iCur];
}
for (const auto& next : m_vNP[iCur])
{
if (iParent == next)
{
continue;
}
dfs2(next, iCur);
}
}
vector m_vChildDis;//距离子孙节点之和
vector m_vChildNum;//子孙数量之和+1
vector m_vTotalNum;
int m_n;
vector < vector > m_vNP;
};

2023年 6月

class Solution {
public:
vector sumOfDistancesInTree(int n, vector<vector>& edges) {
m_vTotalDis.resize(n);
m_vNodeNum.resize(n);
m_vLeve.resize(n);
m_vParent.assign(n, -1);
CNeiBo2 neiBo(n, edges, false);
DSFLeveAndNodeCount(neiBo.m_vNeiB, 0, -1);
m_vTotalDis[0] = std::accumulate(m_vLeve.begin(), m_vLeve.end(), 0);
//必须按父子顺序出来
DFS(neiBo.m_vNeiB, 0, -1);
return m_vTotalDis;
}
void DFS(vector<vector>& neiBo, int iCur, int iParent)
{
if (-1 != iParent)
{
m_vTotalDis[iCur] = m_vTotalDis[m_vParent[iCur]] - m_vNodeNum[iCur] + (m_vNodeNum[0] - m_vNodeNum[iCur]);
}
for (const auto& next : neiBo[iCur])
{
if (next == iParent)
{
continue;
}
DFS(neiBo, next, iCur);
}
}
void DSFLeveAndNodeCount( vector<vector>& neiBo,int iCur,int iParent)
{
if (-1 != iParent)
{
m_vLeve[iCur] = m_vLeve[iParent] + 1;
m_vParent[iCur] = iParent;
}
m_vNodeNum[iCur] = 1;
for (const auto& next : neiBo[iCur])
{
if (next == iParent)
{
continue;
}
DSFLeveAndNodeCount(neiBo,next, iCur);
m_vNodeNum[iCur] += m_vNodeNum[next];
}
}
vector m_vTotalDis,m_vNodeNum,m_vLeve,m_vParent;
};