Count Pairs Of Nodes

Hard

#1631Time: O(E + Q * n^2), where `E` is the number of edges, `Q` is the number of queries, and `n` is the number of nodes. The `O(E)` part is for pre-computation, and `O(Q * n^2)` is for processing the queries.Space: O(n + E), where `n` is the number of nodes and `E` is the number of edges. This is for storing the `degrees` array and the `common_edges` map.

Patterns

Two Pointers

Algorithms

Binary Search Sorting

Data structures

Array Graph

Prompt

You are given an undirected graph defined by an integer n, the number of nodes, and a 2D integer array edges, the edges in the graph, where edges[i] = [u_i, v_i] indicates that there is an undirected edge between u_i and v_i. You are also given an integer array queries.

Let incident(a, b) be defined as the number of edges that are connected to either node a or b.

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

a < b
incident(a, b) > queries[j]

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

Note that there can be multiple edges between the same two nodes.

Example 1:

Input: n = 4, edges = [[1,2],[2,4],[1,3],[2,3],[2,1]], queries = [2,3]
Output: [6,5]
Explanation: The calculations for incident(a, b) are shown in the table above.
The answers for each of the queries are as follows:
- answers[0] = 6. All the pairs have an incident(a, b) value greater than 2.
- answers[1] = 5. All the pairs except (3, 4) have an incident(a, b) value greater than 3.

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⁴
1 <= edges.length <= 10⁵
1 <= u_i, v_i <= n
u_i!= v_i
1 <= queries.length <= 20
0 <= queries[j] < edges.length

Approaches

2 approaches with complexity analysis and trade-offs.

The most direct way to solve the problem is to check every possible pair of nodes (a, b) for each query. This involves calculating the incident(a, b) value for each pair and comparing it with the query value.

Algorithm

First, pre-process the edges to calculate the degree of each node and the number of common edges between any two nodes.
We can use an array degrees of size n+1 to store the degree of each node.
We can use a hash map, say common_edges, to store the number of edges between any pair of nodes (u, v). The key can be a combined value like u * (n+1L) + v (assuming u < v) to uniquely identify the pair.
This pre-computation takes O(E) time, where E is the number of edges.
Then, for each query q in queries:
- We initialize a counter to zero.
- We iterate through all possible pairs of distinct nodes (a, b) where a < b.
- For each pair, we calculate incident(a, b) = degrees[a] + degrees[b] - common_edges.getOrDefault((a,b), 0).
- If incident(a, b) > q, we increment the counter.
- After checking all pairs, the counter holds the answer for the query q.
The number of pairs is n * (n-1) / 2, which is O(n^2). Since we do this for each query, this part takes O(Q * n^2).

Walkthrough

This approach begins by pre-calculating two key pieces of information: the degree of each node and the number of common edges for every pair of nodes. The degrees can be stored in an array, and the common edge counts in a hash map for efficient lookup. This setup phase takes time proportional to the number of edges, O(E).

Once the pre-computation is done, we process each query one by one. For a given query q, we iterate through every unique pair of nodes (a, b) with a < b. For each pair, we compute incident(a, b) using the pre-calculated degrees and common edge count. The formula is incident(a, b) = degree(a) + degree(b) - common_edges(a, b). If this value is greater than q, we increment a counter for the current query. This process is repeated for all O(n^2) pairs.

1class Solution {2    public int[] countPairs(int n, int[][] edges, int[] queries) {3        int[] degrees = new int[n + 1];4        Map<Long, Integer> commonEdges = new HashMap<>();5        for (int[] edge : edges) {6            int u = edge[0];7            int v = edge[1];8            degrees[u]++;9            degrees[v]++;10            if (u > v) {11                int temp = u;12                u = v;13                v = temp;14            }15            long key = (long) u * (n + 1) + v;16            commonEdges.put(key, commonEdges.getOrDefault(key, 0) + 1);17        }18 19        int[] answers = new int[queries.length];20        for (int i = 0; i < queries.length; i++) {21            int q = queries[i];22            int count = 0;23            for (int a = 1; a <= n; a++) {24                for (int b = a + 1; b <= n; b++) {25                    long key = (long) a * (n + 1) + b;26                    int common = commonEdges.getOrDefault(key, 0);27                    if (degrees[a] + degrees[b] - common > q) {28                        count++;29                    }30                }31            }32            answers[i] = count;33        }34        return answers;35    }36}

Complexity

Time

O(E + Q * n^2), where `E` is the number of edges, `Q` is the number of queries, and `n` is the number of nodes. The `O(E)` part is for pre-computation, and `O(Q * n^2)` is for processing the queries.

Space

O(n + E), where `n` is the number of nodes and `E` is the number of edges. This is for storing the `degrees` array and the `common_edges` map.

Trade-offs

Pros

Simple to understand and implement.
It correctly solves the problem for small inputs.

Cons

The time complexity is too high for the given constraints on n, leading to a 'Time Limit Exceeded' error on most platforms.

The brute-force approach is too slow due to the O(n^2) loop for each query. We can optimize the counting by reformulating the condition incident(a, b) > q as degree(a) + degree(b) > q + common_edges(a, b). The common_edges term makes it dependent on the specific pair. A better strategy is to first count all pairs (a, b) satisfying the simpler condition degree(a) + degree(b) > q, and then subtract the pairs that were counted incorrectly due to the common_edges term.

Algorithm

Pre-computation:
- Calculate the degrees of all nodes by iterating through edges. This takes O(E).
- Store the count of common_edges for each pair that has at least one edge. A hash map is suitable for this. This also takes O(E).
- Create a copy of the degrees array (for nodes 1 to n) and sort it. Let's call it sorted_degrees. This takes O(n log n).
Query Processing: For each query q:
- Initial Count (Overestimation): Use a two-pointer technique on sorted_degrees to count pairs (a, b) where degree(a) + degree(b) > q. This takes O(n).
  - Initialize left = 0, right = n-1, and count = 0.
  - While left < right:
    - If sorted_degrees[left] + sorted_degrees[right] > q, then all pairs (left, k) for k from left+1 to right are valid. Add right - left to count and decrement right.
    - Otherwise, increment left.
- Correction: The initial count is an overestimation. We must subtract pairs (u, v) that were counted but are invalid. These are pairs where degree(u) + degree(v) > q but degree(u) + degree(v) - common_edges(u, v) <= q.
  - Iterate through the common_edges map. For each entry, check if it satisfies the correction condition. If so, decrement the total count for the query.
  - This step takes O(E_unique), where E_unique is the number of pairs with shared edges (E_unique <= E).
The final answer for the query is the corrected count.

Walkthrough

This approach significantly improves performance by avoiding the O(n^2) iteration. The key insight is to separate the counting process into two main parts: an initial, efficient overestimation, followed by a targeted correction.

First, we pre-calculate node degrees and common edge counts, similar to the brute-force method. Then, we create a sorted array of all node degrees. For each query q, we use this sorted array to find the number of pairs (a, b) where the sum of their degrees degree(a) + degree(b) is greater than q. This can be done in O(n) time using a two-pointer approach.

This initial count, however, doesn't account for the common_edges term. It overcounts by including pairs (a, b) that satisfy degree(a) + degree(b) > q but fail the actual condition degree(a) + degree(b) - common_edges(a, b) > q. To fix this, we perform a correction step. We iterate through only those pairs that have common edges (which we stored in our map). For each such pair, we check if it was wrongly included in our initial count. If degree(a) + degree(b) > q but degree(a) + degree(b) - common_edges(a, b) <= q, we decrement our count. This correction step's complexity depends on the number of unique pairs with edges, which is at most E.

1class Solution {2    public int[] countPairs(int n, int[][] edges, int[] queries) {3        int[] degrees = new int[n + 1];4        Map<Long, Integer> commonEdges = new HashMap<>();5        for (int[] edge : edges) {6            int u = edge[0];7            int v = edge[1];8            degrees[u]++;9            degrees[v]++;10            if (u > v) {11                int temp = u;12                u = v;13                v = temp;14            }15            long key = (long) u * (n + 1) + v;16            commonEdges.put(key, commonEdges.getOrDefault(key, 0) + 1);17        }18 19        int[] sortedDegrees = new int[n];20        for (int i = 0; i < n; i++) {21            sortedDegrees[i] = degrees[i + 1];22        }23        Arrays.sort(sortedDegrees);24 25        int[] answers = new int[queries.length];26        for (int i = 0; i < queries.length; i++) {27            int q = queries[i];28            int totalPairs = 0;29            int left = 0, right = n - 1;30            while (left < right) {31                if (sortedDegrees[left] + sortedDegrees[right] > q) {32                    totalPairs += (right - left);33                    right--;34                } else {35                    left++;36                }37            }38 39            int corrections = 0;40            for (Map.Entry<Long, Integer> entry : commonEdges.entrySet()) {41                long key = entry.getKey();42                int commonCount = entry.getValue();43                int u = (int) (key / (n + 1));44                int v = (int) (key % (n + 1));45                int sumDeg = degrees[u] + degrees[v];46                if (sumDeg > q && sumDeg - commonCount <= q) {47                    corrections++;48                }49            }50            answers[i] = totalPairs - corrections;51        }52        return answers;53    }54}

Complexity

Time

O(E + n log n + Q * (n + E_unique)), where `E_unique` is the number of pairs with at least one edge between them (`E_unique <= E`). The `O(n log n)` is for sorting degrees. Each of the `Q` queries takes `O(n)` for the two-pointer scan and `O(E_unique)` for the correction.

Space

O(n + E), for storing degrees, a sorted copy of degrees, and the common edges map.

Trade-offs

Pros

Significantly more efficient than the brute-force approach.
The two-pointer technique for counting pairs based on degree sum is very fast.
Handles the problem constraints effectively.

Cons

The correction step iterates through all unique edges for every query, which could be slow if both the number of queries and edges are large.

Solutions

Solution

1class Solution {2public3  int[] countPairs(int n, int[][] edges, int[] queries) {4    int[] cnt = new int[n];5    Map<Integer, Integer> g = new HashMap<>();6    for (var e : edges) {7      int a = e[0] - 1, b = e[1] - 1;8      ++cnt[a];9      ++cnt[b];10      int k = Math.min(a, b) * n + Math.max(a, b);11      g.merge(k, 1, Integer : : sum);12    }13    int[] s = cnt.clone();14    Arrays.sort(s);15    int[] ans = new int[queries.length];16    for (int i = 0; i < queries.length; ++i) {17      int t = queries[i];18      for (int j = 0; j < n; ++j) {19        int x = s[j];20        int k = search(s, t - x, j + 1);21        ans[i] += n - k;22      }23      for (var e : g.entrySet()) {24        int a = e.getKey() / n, b = e.getKey() % n;25        int v = e.getValue();26        if (cnt[a] + cnt[b] > t && cnt[a] + cnt[b] - v <= t) {27          --ans[i];28        }29      }30    }31    return ans;32  }33private34  int search(int[] arr, int x, int i) {35    int left = i, right = arr.length;36    while (left < right) {37      int mid = (left + right) >> 1;38      if (arr[mid] > x) {39        right = mid;40      } else {41        left = mid + 1;42      }43    }44    return left;45  }46}

1class Solution {2public:3  vector<int> countPairs(int n, vector<vector<int>> &edges,4                         vector<int> &queries) {5    vector<int> cnt(n);6    unordered_map<int, int> g;7    for (auto &e : edges) {8      int a = e[0] - 1, b = e[1] - 1;9      ++cnt[a];10      ++cnt[b];11      int k = min(a, b) * n + max(a, b);12      ++g[k];13    }14    vector<int> s = cnt;15    sort(s.begin(), s.end());16    vector<int> ans(queries.size());17    for (int i = 0; i < queries.size(); ++i) {18      int t = queries[i];19      for (int j = 0; j < n; ++j) {20        int x = s[j];21        int k = upper_bound(s.begin() + j + 1, s.end(), t - x) - s.begin();22        ans[i] += n - k;23      }24      for (auto &[k, v] : g) {25        int a = k / n, b = k % n;26        if (cnt[a] + cnt[b] > t && cnt[a] + cnt[b] - v <= t) {27          --ans[i];28        }29      }30    }31    return ans;32  }33};

1class Solution:2    def countPairs(self, n: int, edges: List[List[int]], queries: List[int]) -> List[int]: cnt = [0] * n g = defaultdict(int) for a, b in edges: a, b = a - 1, b - 1 a, b = min(a, b), max(a, b) cnt[a] += 1 cnt[b] += 1 g[(a, b)] += 1 s = sorted(cnt) ans = [0] * len(queries) for i, t in enumerate(queries): for j, x in enumerate(s): k = bisect_right(s, t - x, lo=j + 1) ans[i] += n - k for (a, b), v in g . items(): if cnt[a] + cnt[b] > t and cnt[a] + cnt[b] - v <= t: ans[i] -= 1 return ans

Video walkthrough

Newsletter

One sharp idea, every week

System design and interview prep — short enough to finish.

No spam. Unsubscribe anytime.

Practice

Same difficulty — related problems to reinforce the pattern.

← Previous

Previous problem

Next problem