https://leetcode.com/problems/count-pairs-of-nodes

It was quite painful, but in fact not very difficult problem, which I solved just in last minute on contest!

Idea is the following:

  1. Let d be dictionary with frequencies for each node
  2. Let d2 be dictionary with frequencies for each edge

What is asked is for each number Q in queries find how many pairs of numbers (i, j) there is, such that total number of edges incedent to i or j is less than Q. There will be O(n^2) such pairs, and there is no way we can check all of them, it will take too much time. However we can notice, that this table is sparse in some sence: we can have only E edges and we need to use this information.

Let us for each Q find answer in two steps:

  1. Imagine, that we do not have edges at all: so what numbers we have inside our table: Imagine, we have frequencies in d equal to 1, 2, 5, 7. Then what we have is 4x4 table:
  1 2 5 7
1 2 3 6 8
2 3 4 7 9
5 6 7 10 12
7 8 9 12 14

And the quesion now is how many elements in this table are more than Q? We can answer this quesiton, using binary search! Also we need to deal with diagonal elements and keep in mind that we count each pair twice.

On the second stage we remember that in fact we have edges and we need to fix some values, so we iterate over counter d2 and do it.

Complexity: time complexity is O(n log n) for each query: we have n rows and we do binary search for each row. Also we need to preprocess all E edges, so final time complexity is O(n log n * q + E*q). Space complexity is O(E + n)

Further discussion: note, that we do not really need binary search here, for each Q - d_sorted[i] value we can precalculate it: the maximum value can be equal to E, so complexity can be reduced to O(E*q), I will add code later.

from bisect import bisect

class Solution:
    def countPairs(self, n, edges, queries):
        d, d2 = Counter(), Counter()
        for x, y in edges:
            d[x] += 1
            d[y] += 1
            d2[(min(x, y), max(x,y))] += 1
   
        d_sorted = sorted(d.values())
        d_sorted = [0]*(n - len(d_sorted)) + d_sorted

        out = []

        for Q in queries:
            ans = sum(n - bisect(d_sorted, Q - d_sorted[i]) for i in range(n))
            ans -= n - bisect(d_sorted, Q//2)
            ans = ans//2

            for x, y in d2:
                ans += (d[x] + d[y] - d2[x, y]) > Q
                ans -= (d[x] + d[y]) > Q
            out.append(ans)

        return out

Better complexity

Here is solution with 2 Pointers apporach, with O(n log n + Eq) complexity. Here we replace binary search with 2 pointers approach, similar what you can use in all Ksum problems.

class Solution:
    def countPairs(self, n, edges, queries):
        d, d2 = Counter(), Counter()
        for x, y in edges:
            d[x] += 1
            d[y] += 1
            d2[(min(x, y), max(x,y))] += 1
   
        d_sorted = sorted(d.values())
        d_sorted = [0]*(n - len(d_sorted)) + d_sorted

        out = []

        for Q in queries:
            ans, j = 0, n - 1
            for i in range(n):
                while j >= 0 and d_sorted[i] + d_sorted[j] > Q: j -= 1
                ans += n - j - 1
            
            ans -= sum(d[i]*2 > Q for i in range(n))
            ans //= 2

            for x, y in d2:
                ans += (d[x] + d[y] - d2[x, y]) > Q
                ans -= (d[x] + d[y]) > Q
            out.append(ans)

        return out

If you like the solution, you can upvote it on leetcode discussion section: Problem 1782