# 1782. Count Pairs Of Nodes¶

• Time: $O(q(n + |\texttt{edges}|))$
• Space: $O(n + |\texttt{edges}|)$
  1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 class Solution { public: vector countPairs(int n, vector>& edges, vector& queries) { vector ans(queries.size()); // count[i] := the number of edges of node i vector count(n + 1); // shared[i][j] := the number of edges incident to i or j, where i < j vector> shared(n + 1); for (const vector& edge : edges) { const int u = edge[0]; const int v = edge[1]; ++count[u]; ++count[v]; ++shared[min(u, v)][max(u, v)]; } vector sortedCount(count); ranges::sort(sortedCount); int k = 0; for (const int query : queries) { for (int i = 1, j = n; i < j;) if (sortedCount[i] + sortedCount[j] > query) // sortedCount[i] + sortedCount[j] > query // sortedCount[i + 1] + sortedCount[j] > query // ... // sortedCount[j - 1] + sortedCount[j] > query // So, there are (j - 1) - i + 1 = j - i pairs > query ans[k] += (j--) - i; else ++i; for (int i = 1; i <= n; ++i) for (const auto& [j, sh] : shared[i]) if (count[i] + count[j] > query && count[i] + count[j] - sh <= query) --ans[k]; ++k; } return ans; } }; 
  1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 class Solution { public int[] countPairs(int n, int[][] edges, int[] queries) { int[] ans = new int[queries.length]; // count[i] := the number of edges of node i int[] count = new int[n + 1]; // shared[i][j] := the number of edges incident to i or j, where i < j Map[] shared = new Map[n + 1]; for (int i = 1; i <= n; ++i) shared[i] = new HashMap<>(); for (int[] edge : edges) { final int u = edge[0]; final int v = edge[1]; ++count[u]; ++count[v]; shared[Math.min(u, v)].merge(Math.max(u, v), 1, Integer::sum); } int[] sortedCount = count.clone(); Arrays.sort(sortedCount); int k = 0; for (final int query : queries) { for (int i = 1, j = n; i < j;) if (sortedCount[i] + sortedCount[j] > query) // sortedCount[i] + sortedCount[j] > query // sortedCount[i + 1] + sortedCount[j] > query // ... // sortedCount[j - 1] + sortedCount[j] > query // So, there are (j - 1) - i + 1 = j - i pairs > query ans[k] += (j--) - i; else ++i; for (int i = 1; i <= n; ++i) for (Map.Entry p : shared[i].entrySet()) { final int j = p.getKey(); final int sh = p.getValue(); if (count[i] + count[j] > query && count[i] + count[j] - sh <= query) --ans[k]; } ++k; } return ans; } } 
  1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 class Solution: def countPairs(self, n: int, edges: List[List[int]], queries: List[int]) -> List[int]: ans = [0] * len(queries) # count[i] := the number of edges of node i count = [0] * (n + 1) # shared[i][j] := the number of edges incident to i or j, where i < j shared = [collections.Counter() for _ in range(n + 1)] for u, v in edges: count[u] += 1 count[v] += 1 shared[min(u, v)][max(u, v)] += 1 sortedCount = sorted(count) for k, query in enumerate(queries): i = 1 j = n while i < j: if sortedCount[i] + sortedCount[j] > query: # sortedCount[i] + sortedCount[j] > query # sortedCount[i + 1] + sortedCount[j] > query # ... # sortedCount[j - 1] + sortedCount[j] > query # So, there are (j - 1) - i + 1 = j - i pairs > query ans[k] += j - i j -= 1 else: i += 1 for i in range(1, n + 1): for j, sh in shared[i].items(): if count[i] + count[j] > query and count[i] + count[j] - sh <= query: ans[k] -= 1 return ans