1782. Count Pairs Of Nodes

• Time: $O(q(n + |\texttt{edges}|))$
• Space: $O(n + |\texttt{edges}|)$

C++JavaPython

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

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

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