Count Pairs Of Nodes LeetCode solution in Python

1
class Solution:
2
    def countPairs(
3
        self, n: int, edges: List[List[int]], queries: List[int]
4
    ) -> List[int]:
5
        p_c = [0] * (n + 1)  # point counter
6
        e_c = defaultdict(int)  # edge counter
7

8
        for a, b in edges:
9
            p_c[a] += 1
10
            p_c[b] += 1
11
            if a < b:
12
                e_c[(a, b)] += 1
13
            else:
14
                e_c[(b, a)] += 1
15

16
        f_c = defaultdict(
17
            int
18
        )  # frequency of point counter (compress the individual calculations into massive one)
19
        for i in range(1, n + 1):
20
            f_c[p_c[i]] += 1
21

22
        cnt_amt = defaultdict(
23
            int
24
        )  # frequency of same cnt on pairs without edge adjustment
25
        f_key = list(f_c.keys())  # frequency key
26
        l_f = len(f_key)
27
        for i in range(l_f):
28
            f1 = f_key[i]
29
            cnt_amt[f1 * 2] += (
30
                f_c[f1] * (f_c[f1] - 1) // 2
31
            )  # 3 points with 4, then you get 8
32
            # for a frequency of 3 * (3-2) // 2 = 3
33
            for j in range(i + 1, l_f):
34
                f2 = f_key[j]
35
                cnt_amt[f1 + f2] += (
36
                    f_c[f1] * f_c[f2]
37
                )  # 2 points with 2, 3 points with 4
38
                # will have 6 with frequency of 2 * 3 = 6
39

40
        for (u, v), z in e_c.items():  # edge adjustment
41
            s = p_c[u] + p_c[v]
42
            t = s - z
43
            cnt_amt[s] -= 1
44
            cnt_amt[t] += 1
45

46
        cnt_keys = sorted(
47
            cnt_amt.keys(), key=lambda x: -x
48
        )  # accumulate the frequency by descending order
49
        res = defaultdict(int)
50
        tmp = 0
51
        for key in cnt_keys:
52
            if cnt_amt[key]:  # if it is 0 then dont bother to keep them
53
                tmp += cnt_amt[key]
54
                res[key] = tmp
55

56
        res_keys = sorted(res.keys())  # bisection search
57
        for i, q in enumerate(queries):
58
            local = bisect.bisect(res_keys, q)
59
            try:
60
                queries[i] = res[res_keys[local]]
61
            except:
62
                queries[i] = (
63
                    0  # some test cases are annoying by requesting an unreasonable large number
64
                )
65

66
        return queries