Leetcode 2421. Number of Good Paths #179
Description
这道题是类似多叉树中求path的问题。
我总结了下，对于树/图问题，我的解决方法有如下几种：
- 直接DFS/BFS(对于二维数组移动问题比较常见)
- 树形DP问题(比如Binary Tree Camera / Maximum path sum of the tree)，可解决多叉树，但关键是要给定起始点
- 并查集算法(针对图)
所以我一开始的想法，看到是Path问题，就想利用树状DP来解决：
- 找到根节点
- 然后利用后序遍历，返回Map解决问题，注意相邻问题
- 循环中，利用相乘计算ans
这里有一些问题：
- 首先这类问题根本没有所谓的根节点，因为每个点都可以叫做根节点。除非题目指定root或某个vertex。比如1443. Minimum Time to Collect All Apples in a Tree
- 在递归函数中计算ans时，在边多的情况下不可避免地会超时
class Solution {
int ans = 0;
public int numberOfGoodPaths(int[] vals, int[][] edges) {
// key: the index of the node, value: the value of the node
Map<Integer, Integer> nodeValueMap = new HashMap<>();
int n = vals.length;
ans = n;
for (int i = 0; i < vals.length; i++) {
nodeValueMap.put(i, vals[i]);
}
// convert edges to graph
Map<Integer, List<Integer>> graph = new HashMap<>();
for (int[] edge : edges) {
int from = edge[0];
int to = edge[1];
graph.computeIfAbsent(from, k -> new ArrayList<>()).add(to);
graph.computeIfAbsent(to, k -> new ArrayList<>()).add(from);
}
// find root of the multi-tree
// int root = 0;
// for (int i = 0; i < n; i++) {
// int nodeCnt = getNodeCnt(i, graph, new HashSet<>());
// if (nodeCnt == n) {
// root = i;
// break;
// }
// }
// dfs
dfs(graph, -1, 0, nodeValueMap);
return ans;
}
private Map<Integer, Integer> dfs(Map<Integer, List<Integer>> graph, int parent, int cur, Map<Integer, Integer> nodeValueMap) {
Map<Integer, Integer> curNodeValueMap = new HashMap<>();
int val = nodeValueMap.get(cur);
Map<Integer, Integer> cntMap = new HashMap<>();
List<Map<Integer, Integer>> childrenNodeValueMapList = new ArrayList<>();
curNodeValueMap.put(val, 1);
if (graph.containsKey(cur)) {
for (int next : graph.get(cur)) {
if (next == parent) {
continue;
}
Map<Integer, Integer> nextNodeValueMap = dfs(graph, cur, next, nodeValueMap);
childrenNodeValueMapList.add(nextNodeValueMap);
for (int key : nextNodeValueMap.keySet()) {
if (key >= val) {
curNodeValueMap.put(key, curNodeValueMap.getOrDefault(key, 0) + nextNodeValueMap.get(key));
cntMap.put(key, cntMap.getOrDefault(key, 0) + nextNodeValueMap.get(key));
}
if (key == val) {
ans += nextNodeValueMap.get(key);
}
}
}
}
// update ans
for (Map<Integer, Integer> childNodeValueMap : childrenNodeValueMapList) {
for (int key : childNodeValueMap.keySet()) {
if (key >= val) {
if (cntMap.containsKey(key) && cntMap.get(key) >= childNodeValueMap.get(key)) {
cntMap.put(key, cntMap.get(key) - childNodeValueMap.get(key));
}
ans += (childNodeValueMap.get(key) * cntMap.getOrDefault(key, 0));
}
}
}
return curNodeValueMap;
}
private int getNodeCnt(int node, Map<Integer, List<Integer>> graph, Set<Integer> visited) {
visited.add(node);
int cnt = 1;
if (!graph.containsKey(node)) {
return cnt;
}
for (int next : graph.get(node)) {
if (!visited.contains(next)) {
cnt += getNodeCnt(next, graph, visited);
}
}
return cnt;
}
}
感觉代码思路没什么问题，但卡在最后十几个test用例上，最后TLE了。
那么正解是什么呢？如何只用O(n)或者O(nlgn)的时间复杂度呢？
可以这么想，题目要求找到一个起始点value相同的节点，然后在路径中还要要求所有节点的value小于等于起始点value。是不是能将树/图分割成一个个群？即使用并查集算法，从每一个value相同的初始孤立的节点开始；然后问题是如果在并查集中如何merge节点呢？先从value小的开始，按顺序到大的，这样我们就只需要merge相邻节点即可；最后记录所有节点的群，及群的节点个数，求组成的path个数，遍历出现value的节点，用Map存储，key表示群的根节点，val表示同时为根节点相等的value相等的个数；最后我们知道每个群有多少个相同的value节点，然后用数学方法求子集，即count*(count+1)/2。
class Solution {
public int numberOfGoodPaths(int[] vals, int[][] edges) {
// convert edges to graph
Map<Integer, List<Integer>> graph = new HashMap<>();
for (int[] edge : edges) {
int from = edge[0];
int to = edge[1];
graph.computeIfAbsent(from, k -> new ArrayList<>()).add(to);
graph.computeIfAbsent(to, k -> new ArrayList<>()).add(from);
}
// valueMap, key: value, value: index
// ??????? why treemap?? —— connect
TreeMap<Integer, List<Integer>> valueMap = new TreeMap<>();
for (int i = 0; i < vals.length; i++) {
valueMap.computeIfAbsent(vals[i], k -> new ArrayList<>()).add(i);
}
// union find
UnionFind uf = new UnionFind(vals.length);
int ans = 0;
for (List<Integer> list : valueMap.values()) {
Map<Integer, Integer> graphMap = new HashMap<>();
for (int idx : list) {
if (!graph.containsKey(idx)) {
continue;
}
for (int next : graph.get(idx)) {
if (vals[next] <= vals[idx]) {
uf.union_set(idx, next);
}
}
}
for (int idx : list) {
int root = uf.find(idx);
graphMap.put(root, graphMap.getOrDefault(root, 0) + 1);
}
for (int count : graphMap.values()) {
ans += count * (count + 1) / 2;
}
}
return ans;
}
}
class UnionFind {
int[] parent;
int[] rank;
public UnionFind(int size) {
parent = new int[size];
for (int i = 0; i < size; i++)
parent[i] = i;
rank = new int[size];
}
public int find(int x) {
if (parent[x] != x)
parent[x] = find(parent[x]);
return parent[x];
}
public void union_set(int x, int y) {
int xRoot = find(x);
int yRoot = find(y);
if (xRoot == yRoot)
return;
if (rank[xRoot] < rank[yRoot]) {
parent[xRoot] = yRoot;
rank[yRoot] += rank[xRoot];
} else {
parent[yRoot] = xRoot;
rank[xRoot] += rank[yRoot];
}
}
}
Tips:
- 一定要注意遍历下一批节点时，一定要先用containsKey判断下
- 这题还是很难想到用并查集算法的，但很有趣
- 最后的
int count : graphMap.values()，key不重要了，我们只需要计算在一个群中相同value节点的个数
同时，注意多叉树的path问题，可以用排序+并查集解决。
类似题目：