[leetcode] Number of Good Paths

https://leetcode.com/problems/number-of-good-paths/

Number of Good Paths - LeetCode

Approach

1. 시작(혹은 끝)노드 값보다 작은 노드를 거쳐야 path로 인정되기 때문에 노드의 값이 더 높은 노드를 루트 노드로 본다.

2. 값이 같은 노드를 발견하면 path 를 합산한 뒤 하나의 노드로 합친다. 이 과정에서 같은 노드 값이 몇 개가 있는지 알아야 함.

3. 값이 작은 노드들부터 루트 노드를 찾고 path 계산을 해야 한다

 

Code

class Solution {
    vector<int> root;
    int find(int node) {
        if(root[node] == node) return node;
        return root[node] = find(root[node]);
    }
public:
    int numberOfGoodPaths(vector<int>& vals, vector<vector<int>>& edges) {
        int ans = vals.size();
        root.resize(vals.size());
        vector<int> v(vals.size());
        for(int i = 0; i < vals.size(); i++) {
            v[i] = 1;
            root[i] = i;
        }

		// vals 를 기준으로 정렬
        sort(edges.begin(), edges.end(), [&](vector<int>& a1, vector<int>& a2) {
            return max(vals[a1[0]], vals[a1[1]]) < max(vals[a2[0]], vals[a2[1]]);
        });

        for(auto& e : edges) {
            int a = find(e[0]);
            int b = find(e[1]);
            if(vals[a] != vals[b]) {
                if(vals[a] < vals[b]) root[a] = b;
                else root[b] = a;
            }
            else {
                ans += v[a] * v[b];
                root[b] = a;
                v[a] += v[b];
            }
        }

        return ans;
    }
};