【树形dp 换根法 BFS】2581. 统计可能的树根数目

本文涉及知识点

C++BFS算法
动态规划汇总
图论知识汇总
树形dp 换根法 BFS

LeetCode 2581. 统计可能的树根数目

Alice 有一棵 n 个节点的树，节点编号为 0 到 n - 1 。树用一个长度为 n - 1 的二维整数数组 edges 表示，其中 edges[i] = [ai, bi] ，表示树中节点 ai 和 bi 之间有一条边。
Alice 想要 Bob 找到这棵树的根。她允许 Bob 对这棵树进行若干次猜测。每一次猜测，Bob 做如下事情：
选择两个不相等的整数 u 和 v ，且树中必须存在边 [u, v] 。
Bob 猜测树中 u 是 v 的父节点。
Bob 的猜测用二维整数数组 guesses 表示，其中 guesses[j] = [uj, vj] 表示 Bob 猜 uj 是 vj 的父节点。
Alice 非常懒，她不想逐个回答 Bob 的猜测，只告诉 Bob 这些猜测里面至少有 k 个猜测的结果为 true 。
给你二维整数数组 edges ，Bob 的所有猜测和整数 k ，请你返回可能成为树根的节点数目。如果没有这样的树，则返回 0。
示例 1：
输入：edges = [[0,1],[1,2],[1,3],[4,2]], guesses = [[1,3],[0,1],[1,0],[2,4]], k = 3
输出：3
解释：
根为节点 0 ，正确的猜测为 [1,3], [0,1], [2,4]
根为节点 1 ，正确的猜测为 [1,3], [1,0], [2,4]
根为节点 2 ，正确的猜测为 [1,3], [1,0], [2,4]
根为节点 3 ，正确的猜测为 [1,0], [2,4]
根为节点 4 ，正确的猜测为 [1,3], [1,0]
节点 0 ，1 或 2 为根时，可以得到 3 个正确的猜测。
示例 2：
输入：edges = [[0,1],[1,2],[2,3],[3,4]], guesses = [[1,0],[3,4],[2,1],[3,2]], k = 1
输出：5
解释：
根为节点 0 ，正确的猜测为 [3,4]
根为节点 1 ，正确的猜测为 [1,0], [3,4]
根为节点 2 ，正确的猜测为 [1,0], [2,1], [3,4]
根为节点 3 ，正确的猜测为 [1,0], [2,1], [3,2], [3,4]
根为节点 4 ，正确的猜测为 [1,0], [2,1], [3,2]
任何节点为根，都至少有 1 个正确的猜测。
提示：
edges.length == n - 1
2 <= n <= 10⁵
1 <= guesses.length <= 10⁵
0 <= ai, bi, uj, vj <= n - 1
ai != bi
uj != vj
edges 表示一棵有效的树。
guesses[j] 是树中的一条边。
guesses 是唯一的。
0 <= k <= guesses.length

换根法

某棵有根树，根为root，某个儿子为child，则将根从root换成child后，除 $root\leftrightarrow child$ 这条的边的父子关系发生变化外，其它都不边。
mGuesss[x] 记录猜测次数:x=Mask(r,v) = u*n+v
x1 = Mask(root,child)
x2 = Mask(child,root)
则 dp[child] = dp[root] - mGuesss[x1] + mGuress[x2]
分三步：
一，令root为0,计算m_dp[0]。
二，dfs各节点计算m_dp[cur]。
三，统计m_dp中为k的元素数量。

动态规划的状态表示

m_dp[cur]表示以cur为根据猜对父子关系的数量。
空间复杂度: O(n)

动态规划的转移方程

dp[child] = dp[root] - mGuesss[x1] + mGuress[x2]
单个状态的转移方程时间复杂度：O(1) 总时间复杂度：O(n)

动态规划的初始值

dp[0]先计算

动态规划的填表顺序

深度优先，广度优先也可以。

动态规划的返回值

cout(dp.being(),dp.end(),k)

代码（超时)

核心代码

class CNeiBo
{
public:	static vector<vector<int>> Two(int n, vector<vector<int>>& edges, bool bDirect, int iBase = 0) {vector<vector<int>>  vNeiBo(n);for (const auto& v : edges){vNeiBo[v[0] - iBase].emplace_back(v[1] - iBase);if (!bDirect){vNeiBo[v[1] - iBase].emplace_back(v[0] - iBase);}}return vNeiBo;}	static vector<vector<std::pair<int, int>>> Three(int n, vector<vector<int>>& edges, bool bDirect, int iBase = 0){vector<vector<std::pair<int, int>>> vNeiBo(n);for (const auto& v : edges){vNeiBo[v[0] - iBase].emplace_back(v[1] - iBase, v[2]);if (!bDirect){vNeiBo[v[1] - iBase].emplace_back(v[0] - iBase, v[2]);}}return vNeiBo;}static vector<vector<int>> Grid(int rCount, int cCount, std::function<bool(int, int)> funVilidCur, std::function<bool(int, int)> funVilidNext){vector<vector<int>> vNeiBo(rCount * cCount);auto Move = [&](int preR, int preC, int r, int c){if ((r < 0) || (r >= rCount)){return;}if ((c < 0) || (c >= cCount)){return;}if (funVilidCur(preR, preC) && funVilidNext(r, c)){vNeiBo[cCount * preR + preC].emplace_back(r * cCount + c);}};for (int r = 0; r < rCount; r++){for (int c = 0; c < cCount; c++){Move(r, c, r + 1, c);Move(r, c, r - 1, c);Move(r, c, r, c + 1);Move(r, c, r, c - 1);}}return vNeiBo;}static vector<vector<int>> Mat(vector<vector<int>>& neiBoMat){vector<vector<int>> neiBo(neiBoMat.size());for (int i = 0; i < neiBoMat.size(); i++){for (int j = i + 1; j < neiBoMat.size(); j++){if (neiBoMat[i][j]){neiBo[i].emplace_back(j);neiBo[j].emplace_back(i);}}}return neiBo;}
};class Solution {
public:int rootCount(vector<vector<int>>& edges, vector<vector<int>>& guesses, int k) {m_c = edges.size() + 1;m_dp.resize(m_c);m_vNeiBo = CNeiBo::Two(m_c, edges, false);for (const auto& v : guesses) {m_mGuess[Mask(v[0], v[1])]++;}m_dp[0] = DFS1(0, -1);DFS2(0, -1);return count_if(m_dp.begin(), m_dp.end(), [&](int i) {return i >= k; });}int DFS1(int cur, int par) {int ret = 0;if (-1 != par) { ret += m_mGuess[Mask(par, cur)];}for (const auto& next : m_vNeiBo[cur]) {if (next == par) { continue; }ret += DFS1(next, cur);}return ret;}void DFS2(int cur, int par) {if (-1 != par) {m_dp[cur] = m_dp[par];m_dp[cur] -= m_mGuess[Mask(par, cur)];m_dp[cur] += m_mGuess[Mask(cur, par)];}	for (const auto& next : m_vNeiBo[cur]) {if (next == par) { continue; }DFS2(next, cur);}}long long Mask(long long par, int cur) { return m_c * par + cur; }int m_c;vector<int> m_dp;unordered_map<long long, int> m_mGuess;vector < vector <int>> m_vNeiBo;
};

单元测试

template<class T1, class T2>
void AssertEx(const T1& t1, const T2& t2)
{Assert::AreEqual(t1, t2);
}template<class T>
void AssertEx(const vector<T>& v1, const vector<T>& v2)
{Assert::AreEqual(v1.size(), v2.size());for (int i = 0; i < v1.size(); i++){Assert::AreEqual(v1[i], v2[i]);}
}template<class T>
void AssertV2(vector<vector<T>> vv1, vector<vector<T>> vv2)
{sort(vv1.begin(), vv1.end());sort(vv2.begin(), vv2.end());Assert::AreEqual(vv1.size(), vv2.size());for (int i = 0; i < vv1.size(); i++){AssertEx(vv1[i], vv2[i]);}
}namespace UnitTest
{vector<vector<int>> edges, guesses;int k;TEST_CLASS(UnitTest){public:TEST_METHOD(TestMethod0){edges = { {0,1},{1,2},{1,3},{4,2} }, guesses = { {1,3},{0,1},{1,0},{2,4} }, k = 3;auto res = Solution().rootCount(edges, guesses, k);AssertEx(3, res);}TEST_METHOD(TestMethod1){edges = { {0,1},{1,2},{2,3},{3,4} }, guesses = { {1,0},{3,4},{2,1},{3,2} }, k = 1;auto res = Solution().rootCount(edges, guesses, k);AssertEx(5, res);}TEST_METHOD(TestMethod2){edges ={ {1,0},{2,1},{2,3},{4,0},{5,2},{6,1},{0,7},{1,8},{9,6},{10,4},{11,10},{12,8},{8,13},{14,4},{15,9},{9,16},{3,17},{4,18},{6,19},{20,13},{21,20},{19,22},{23,3},{24,0},{25,14},{17,26},{27,3},{3,28},{29,3},{4,30},{31,9},{0,32},{33,12},{34,14},{27,35},{35,36},{37,33},{38,18},{6,39} };guesses ={ {13,8},{4,18},{37,33},{4,30},{1,8},{3,17},{25,14},{0,1},{27,35},{21,20},{6,1},{26,17},{1,2},{8,13},{22,19},{30,4},{4,0},{2,5},{14,4},{9,6},{19,22},{16,9},{5,2},{29,3},{34,14},{8,1},{11,10},{15,9},{10,4},{35,27},{3,27},{33,12},{14,34},{32,0},{14,25},{39,6},{7,0},{4,10},{0,32},{23,3},{20,21},{24,0},{0,7},{1,0},{3,28},{6,9},{8,12},{18,4},{1,6},{2,1},{2,3},{3,29},{9,16},{17,26},{35,36},{13,20},{10,11},{18,38},{3,23},{0,24},{33,37},{12,33},{3,2},{20,13},{17,3} };k =	29;auto res = Solution().rootCount(edges, guesses, k);AssertEx(40, res);}};
}

DFS非常容易超时

DFS稍稍复杂，leetcode就容易超时。
所以：
一，计算出临接表。
二，DFS各节点层次。
三，计算出各节点的孩子。
四，BFS各节点。由于每个节点顶多一个父亲，所以无需判断节点是否重复访问。

class CNeiBo
{
public:	static vector<vector<int>> Two(int n, vector<vector<int>>& edges, bool bDirect, int iBase = 0) {vector<vector<int>>  vNeiBo(n);for (const auto& v : edges){vNeiBo[v[0] - iBase].emplace_back(v[1] - iBase);if (!bDirect){vNeiBo[v[1] - iBase].emplace_back(v[0] - iBase);}}return vNeiBo;}	static vector<vector<std::pair<int, int>>> Three(int n, vector<vector<int>>& edges, bool bDirect, int iBase = 0){vector<vector<std::pair<int, int>>> vNeiBo(n);for (const auto& v : edges){vNeiBo[v[0] - iBase].emplace_back(v[1] - iBase, v[2]);if (!bDirect){vNeiBo[v[1] - iBase].emplace_back(v[0] - iBase, v[2]);}}return vNeiBo;}static vector<vector<int>> Grid(int rCount, int cCount, std::function<bool(int, int)> funVilidCur, std::function<bool(int, int)> funVilidNext){vector<vector<int>> vNeiBo(rCount * cCount);auto Move = [&](int preR, int preC, int r, int c){if ((r < 0) || (r >= rCount)){return;}if ((c < 0) || (c >= cCount)){return;}if (funVilidCur(preR, preC) && funVilidNext(r, c)){vNeiBo[cCount * preR + preC].emplace_back(r * cCount + c);}};for (int r = 0; r < rCount; r++){for (int c = 0; c < cCount; c++){Move(r, c, r + 1, c);Move(r, c, r - 1, c);Move(r, c, r, c + 1);Move(r, c, r, c - 1);}}return vNeiBo;}static vector<vector<int>> Mat(vector<vector<int>>& neiBoMat){vector<vector<int>> neiBo(neiBoMat.size());for (int i = 0; i < neiBoMat.size(); i++){for (int j = i + 1; j < neiBoMat.size(); j++){if (neiBoMat[i][j]){neiBo[i].emplace_back(j);neiBo[j].emplace_back(i);}}}return neiBo;}
};class CDFSLeveChild
{
public:CDFSLeveChild(const vector<vector <int>>& vNeiBo,int root=0):m_vNeiBo(vNeiBo), Leve(m_vLeve){m_vLeve.resize(m_vNeiBo.size());DFS(root, -1);};const vector<int>& Leve;vector<vector<int>> Child() const{vector<vector <int>> vChild(m_vNeiBo.size());for (int i = 0; i < m_vNeiBo.size(); i++) {for (const auto& next : m_vNeiBo[i]) {if (m_vLeve[next] < m_vLeve[i]) { continue; }vChild[i].emplace_back(next);}}return vChild;}
protected:void DFS(int cur, int par) {if (-1 != par) { m_vLeve[cur] = m_vLeve[par] + 1; }for (const auto& next : m_vNeiBo[cur]) {if (next == par) { continue; }DFS(next, cur);}}vector<int> m_vLeve;const vector<vector <int>>& m_vNeiBo;
};
class Solution {
public:int rootCount(vector<vector<int>>& edges, vector<vector<int>>& guesses, int k) {m_c = edges.size() + 1;m_dp.resize(m_c);		m_vNeiBo = CNeiBo::Two(m_c, edges, false);		auto vChilds = CDFSLeveChild(m_vNeiBo).Child();for (const auto& v : guesses) {m_mGuess[Mask(v[0], v[1])]++;}for (int par = 0; par < m_c; par++) {for (int& child : vChilds[par]) {m_dp[0] += m_mGuess[Mask(par, child)];}}	queue<int> que;		que.emplace(0);while (que.size()) {int cur = que.front();que.pop();for (const auto& child : vChilds[cur]) {m_dp[child] = m_dp[cur];m_dp[child] -= m_mGuess[Mask(cur, child)];m_dp[child] += m_mGuess[Mask(child, cur)];que.emplace(child);}}return count_if(m_dp.begin(), m_dp.end(), [&](int i) {return i >= k; });}long long Mask(long long par, int cur) { return m_c * par + cur; }int m_c;vector<int> m_dp;unordered_map<long long, int> m_mGuess;vector < vector <int>> m_vNeiBo;	
};

进一步优化

可以用数组代码映射，算法方向，总共2n-2条边。假定根为0的树。
如果这条边是子节点执行父节点，则此边数是child。如果方向相反则是n + child。
运行速度大约提高了20%。

class Solution {
public:int rootCount(vector<vector<int>>& edges, vector<vector<int>>& guesses, int k) {m_c = edges.size() + 1;m_dp.resize(m_c);	vector<int> vGuess(m_c * 2);m_vNeiBo = CNeiBo::Two(m_c, edges, false);	CDFSLeveChild dfs(m_vNeiBo);auto vChilds = dfs.Child();auto Mask = [&](int par, int child) {if (dfs.Leve[par] < dfs.Leve[child]) {return child;}return par + m_c;};for (const auto& v : guesses) {vGuess[Mask(v[0], v[1])]++;}for (int par = 0; par < m_c; par++) {for (int& child : vChilds[par]) {m_dp[0] += vGuess[Mask(par, child)];}}	queue<int> que;		que.emplace(0);while (que.size()) {int cur = que.front();que.pop();for (const auto& child : vChilds[cur]) {m_dp[child] = m_dp[cur];m_dp[child] -= vGuess[Mask(cur, child)];m_dp[child] += vGuess[Mask(child, cur)];que.emplace(child);}}return count_if(m_dp.begin(), m_dp.end(), [&](int i) {return i >= k; });}int m_c;vector<int> m_dp;	vector < vector <int>> m_vNeiBo;	
};

DFS序+差分数组

root和它的某个后代childchild换根。则到 $childchild\leftrightarrow root$ 这条路径上的边都反转。可以用差分数组。
childchild和它的祖先不是连续的，但他们的DFS序是连续的。
此方案不好理解，实现也不简单。备用。

扩展阅读

视频课程

先学简单的课程，请移步CSDN学院，听白银讲师（也就是鄙人）的讲解。
https://edu.csdn.net/course/detail/38771

如何你想快速形成战斗了，为老板分忧，请学习C#入职培训、C++入职培训等课程
https://edu.csdn.net/lecturer/6176

我想对大家说的话
《喜缺全书算法册》以原理、正确性证明、总结为主。
按类别查阅鄙人的算法文章，请点击《算法与数据汇总》。
有效学习：明确的目标及时的反馈拉伸区（难度合适）专注
闻缺陷则喜(喜缺)是一个美好的愿望，早发现问题，早修改问题，给老板节约钱。
子墨子言之：事无终始，无务多业。也就是我们常说的专业的人做专业的事。
如果程序是一条龙，那算法就是他的是睛