好路径数目-Leetcode6191

好路径数目

题目描述

给你一棵 n 个节点的树（连通无向无环的图），节点编号从 0 到 n - 1 且恰好有 n - 1 条边。

给你一个长度为 n 下标从 0 开始的整数数组 vals ，分别表示每个节点的值。同时给你一个二维整数数组 edges ，其中 edges[i] = [ai, bi] 表示节点 ai 和 bi 之间有一条 无向 边。

一条 好路径 需要满足以下条件：

开始节点和结束节点的值 相同 。
开始节点和结束节点中间的所有节点值都 小于等于 开始节点的值（也就是说开始节点的值应该是路径上所有节点的最大值）。
请你返回不同好路径的数目。

注意，一条路径和它反向的路径算作 同一 路径。比方说， 0 -> 1 与 1 -> 0 视为同一条路径。单个节点也视为一条合法路径。

Example

输入：vals = [1,3,2,1,3], edges = [[0,1],[0,2],[2,3],[2,4]]
输出：6
解释：总共有 5 条单个节点的好路径。
还有 1 条好路径：1 -> 0 -> 2 -> 4 。
（反方向的路径 4 -> 2 -> 0 -> 1 视为跟 1 -> 0 -> 2 -> 4 一样的路径）
注意 0 -> 2 -> 3 不是一条好路径，因为 vals[2] > vals[0] 。

输入：vals = [1,1,2,2,3], edges = [[0,1],[1,2],[2,3],[2,4]]
输出：7
解释：总共有 5 条单个节点的好路径。
还有 2 条好路径：0 -> 1 和 2 -> 3 。

输入：vals = [1], edges = []
输出：1
解释：这棵树只有一个节点，所以只有一条好路径。

Solution

排序+并查集

将节点按权值从小到大排序，用并查集维护点集，从小到大开始处理，对于每个节点，将所有满足：

• 与这个节点相邻
• 权值小于这个节点权值
• 已存在于并查集中

的节点合并到并查集中，并且记录并查集中与当前节点权值相同的点有多少个，由于连通的无向无环图，因此任意两点间只有一条路径（两条则出现环），因此任意两个同权节点间可形成一条好路径（由于点是按照权值从小到大处理，因此此时并查集中的点权值一定小于等于当前点权值，因此路径上点的权值一定小于等于端点权值，因此是好路径），最后将每个节点计算的好路径数量求和，注意每单个节点也算一条好路径。

Code

class Solution {
private:
    int fa[int(3*1e4 + 10)];
public:
    int find(int x) {
        if(fa[x] == -1) return -1;
        if(fa[x] == x) return x;
        return fa[x] = find(fa[x]);
    }

    void merge(int x, int y) {
        int a = find(x);
        int b = find(y);
        if(a == b) return;
        fa[b] = a;
        return;
    }

    int numberOfGoodPaths(vector<int>& vals, vector<vector<int>>& edges) {
        vector<pair<int, int>> nodes;
        unordered_map<int, int> mm[vals.size()];
        for(int i = 0; i < vals.size(); ++i) {
            nodes.push_back(make_pair(vals[i], i));
            mm[i][vals[i]] = 1;
        }
        vector<vector<int>> mp(vals.size());
        for(auto k : edges) {
            mp[k[0]].push_back(k[1]);
            mp[k[1]].push_back(k[0]);
        }
        for(int i = 0; i < vals.size(); ++i) {
            fa[i] = -1;
        }
        sort(nodes.begin(), nodes.end());
        int ans = 0;
        for(int i = 0; i < nodes.size(); ++i) {
            fa[nodes[i].second] = nodes[i].second;
            // cout << i << "----" << endl;
            for(auto k : mp[nodes[i].second]) {
                if(find(k) != -1 && vals[k] <= nodes[i].first) {
                    mm[find(nodes[i].second)][nodes[i].first] += mm[find(k)][nodes[i].first];
                    merge(nodes[i].second, k);
                }
            }
            ans += mm[find(nodes[i].second)][nodes[i].first] - 1;
        }
        return ans + vals.size();
    }
};