Problem: Count Pairs Of Nodes

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]
= [ui, vi]` indicates that there is an undirected edge between `ui` and `vi`.
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 `jth` 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 `jth` query.

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


Python solution:
import math
import itertools
import bisect
import collections
import string
import heapq
import functools
import sortedcontainers
from typing import List, Dict, Tuple, Iterator

class Solution:
  def countPairs(self, n: int, edges: List[List[int]], queries: List[int]) -> List[int]:
    ans = [0] * len(queries)

    count = [0] * (n + 1)

    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:
          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


Java solution:
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<Integer, Integer>[] 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<Integer, Integer> 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;
  }
}


C++ solution:
class Solution {
 public:
  vector<int> countPairs(int n, vector<vector<int>>& edges,
                         vector<int>& queries) {
    vector<int> ans(queries.size());

    // count[i] := the number of edges of node i
    vector<int> count(n + 1);

    // shared[i][j] := the number of edges incident to i or j, where i < j
    vector<unordered_map<int, int>> shared(n + 1);

    for (const vector<int>& edge : edges) {
      const int u = edge[0];
      const int v = edge[1];
      ++count[u];
      ++count[v];
      ++shared[min(u, v)][max(u, v)];
    }

    vector<int> 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;
  }
};