LeetCode 1782. Count Pairs Of Nodes

You are given an undirected graph represented by an integer n, which is the number of nodes, and edges, where edges[i] = [ui, vi] which indicates that there is an undirected edge between ui and vi. You are also given an integer array queries.

The answer to the jth query is the number of pairs of nodes (a, b) that satisfy the following conditions:

• a < b

• cnt is strictly greater than queries[j], where cnt is the number of edges incident to a or b.

Return an array answers such that answers.length == queries.length and answers[j] is the answer of the jth query.

Note that there can be repeated edges.

Example 1:

Input: n = 4, edges = [[1,2],[2,4],[1,3],[2,3],[2,1]], queries = [2,3]
Output: [6,5]
Explanation: The number of edges incident to at least one of each pair is shown above.

Example 2:

Input: n = 5, edges = [[1,5],[1,5],[3,4],[2,5],[1,3],[5,1],[2,3],[2,5]], queries = [1,2,3,4,5]
Output: [10,10,9,8,6]

Constraints:

• 2 <= n <= 2 * 10^4

• 1 <= edges.length <= 10^5

• 1 <= ui, vi <= n

• ui != vi

• 1 <= queries.length <= 20

• 0 <= queries[j] < edges.length

Solution:

English Version in Youtube

中文版解答Youtube Link

中文版解答Bilibili Link

class Solution {
public:
    vector<int> countPairs(int n, vector<vector<int>>& edges, vector<int>& queries) {
        vector<int> sizes(n+1, 0);
        vector<unordered_map<int, int>> graph(n + 1);
        for (const vector<int>& edge : edges) {
            graph[edge[0]][edge[1]]++;
            graph[edge[1]][edge[0]]++;
            sizes[edge[0]]++;
            sizes[edge[1]]++;
        }
        
        vector<int> prefix_sum(100001, 0);
        for (int i = 1; i < sizes.size(); i++) {
            prefix_sum[sizes[i]]++;
        }
        for (int i = 100000; i >= 1; i--) {
            prefix_sum[i-1] += prefix_sum[i];
        }
        
        vector<int> answer;
        for (int query : queries) {
            int total = 0;
            
            for (int i = 1; i <= n; i++) {
                const unordered_map<int, int>& mymap = graph[i];
                
                int target = query - sizes[i];
                int current = 0;
                for (const auto& it : mymap) {
                    int neighbor = it.first;
                    if (sizes[neighbor] > target) {
                        current--;
                    }
                    int count = sizes[i] + sizes[neighbor] - it.second; 
                    if (count > query) {
                        current++;
                    }
                }
                if (sizes[i] > target) {
                    current--;
                }
                current += prefix_sum[target < 0 ? 0 : target+1];
                total += current;
            }
            answer.push_back(total / 2);
        }
        
        return answer;
    }
};