Leetcode 2421: Number of Good Paths

There is a tree (i.e. a connected, undirected graph with no cycles) consisting of n nodes numbered from 0 to n - 1 and exactly n - 1 edges.

You are given a 0-indexed integer array vals of length n where vals[i] denotes the value of the ith node. You are also given a 2D integer array edges where edges[i] = [ai, bi] denotes that there exists an undirected edge connecting nodes ai and bi.

A good path is a simple path that satisfies the following conditions:

1. The starting node and the ending node have the same value.
2. All nodes between the starting node and the ending node have values less than or equal to the starting node (i.e. the starting node’s value should be the maximum value along the path).

Return the number of distinct good paths.

Note that a path and its reverse are counted as the same path. For example, 0 -> 1 is considered to be the same as 1 -> 0. A single node is also considered as a valid path.

Example 1:

Input: vals = [1,3,2,1,3], edges = [[0,1],[0,2],[2,3],[2,4]]
Output: 6
Explanation: There are 5 good paths consisting of a single node.
There is 1 additional good path: 1 -> 0 -> 2 -> 4.
(The reverse path 4 -> 2 -> 0 -> 1 is treated as the same as 1 -> 0 -> 2 -> 4.)
Note that 0 -> 2 -> 3 is not a good path because vals[2] > vals[0].

Example 2:

Input: vals = [1,1,2,2,3], edges = [[0,1],[1,2],[2,3],[2,4]]
Output: 7
Explanation: There are 5 good paths consisting of a single node.
There are 2 additional good paths: 0 -> 1 and 2 -> 3.

Example 3:

Input: vals = [1], edges = []
Output: 1
Explanation: The tree consists of only one node, so there is one good path.

Constraints:

• n == vals.length
• 1 <= n <= 3 * 104
• 0 <= vals[i] <= 105
• edges.length == n - 1
• edges[i].length == 2
• 0 <= ai, bi < n
• ai != bi
• edges represents a valid tree.

Problem Analysis:

At first, we can use depth search first algorithm, we iterate at every node, through the edges, if meeting a node with same value, then it’s a good path. The code is bellow:

class Solution {
public:
    int numberOfGoodPaths(vector<int>& vals, vector<vector<int>>& edges) {
        int res{0};
        unordered_map<int, vector<int>> data;
        for (const auto& item : edges) {
            data[item.front()].push_back(item.back());
            data[item.back()].push_back(item.front());
        }
        vector<vector<bool>> visited(vals.size(), vector<bool>(vals.size(), false));
        
        function<void(int, int, int, unordered_set<int>&)> dfs = [&](int start, int end, int current, unordered_set<int>& indices) {
            if (!visited[start][end]) {
                visited[start][end] = true;
                visited[end][start] = true;
                if (vals[start] == vals[end]) ++res;
            }
            indices.insert(end);
            
            auto it = data.find(end);
            if (it == data.end()) return;
            for (const auto& item : it->second) {
                if (indices.count(item) || vals[item] > vals[start]) continue;
                dfs(start, item, current, indices);
            }
            indices.erase(end);
        };
        unordered_set<int> indices;
        for (int i = 0; i < vals.size(); ++i) dfs(i, i, vals[i], indices);
        
        return res;
    }
};

But the time complexity is O(n²) at the worst case. We need to find better solution.

Here we can collect nodes by value from small to big. We first collect the smallest value nodes. And compute good paths. As if these nodes can connect, there can form another good paths. So every time we can collect the same value nodes and check if they form disjoint sets, in the same disjoint set, every two nodes with same value can form a good path. Then iterate finished, we got the answer.

Solution

class Solution {
public:
    int numberOfGoodPaths(vector<int>& vals, vector<vector<int>>& edges) {
        int res{0};
        vector<vector<int>> adjencency(vals.size(), vector<int>());
        for (const auto& item : edges) {
            if (vals[item.front()] >= vals[item.back()])
                adjencency[item.front()].push_back(item.back());
            if (vals[item.front()] <= vals[item.back()])
                adjencency[item.back()].push_back(item.front());
        }
        map<int, vector<int>> val2nodes;
        for (int i = 0; i < vals.size(); ++i)
            val2nodes[vals[i]].push_back(i);
        
        vector<int> par(vals.size());
        vector<int> rank(vals.size(), 0);
        iota(par.begin(), par.end(), 0);
        function<int(int)> find = [&](int index) {
            if (par[index] == index) return index;
            return par[index] = find(par[index]);
        };
        function<void(int, int)> unite = [&](int lhs, int rhs) {
            lhs = find(lhs);
            rhs = find(rhs);
            if (lhs == rhs) return;
            if (rank[lhs] < rank[rhs]) {
                par[lhs] = rhs;
            } else {
                par[rhs] = lhs;
                if (rank[lhs] == rank[rhs]) ++rank[lhs];
            }
        };
        for (const auto& item : val2nodes) {
            for (const auto& index : item.second) 
                for (const auto& neighbor : adjencency[index])
                    unite(index, neighbor);
            unordered_map<int, int> group2size;
            for (const auto& index : item.second)
                ++group2size[find(index)];
            for (const auto& [groupID, size] : group2size)
                res += size * (size + 1) / 2;
        }
        
        return res;
    }
};

time complexity is O(nodes)

Space complexity is O(nodes)