solution.py

"""
Code generated by https://github.com/goodstudyqaq/leetcode-local-tester
"""
try:
    from utils.python3.help import *
except ImportError:
    pass # In leetcode environment, we don't need to import the help file.

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

        num = [0] * (n + 1)
        M = {}
        f = [0] * (n)
        for i in range(m):
            u, v = edges[i]
            if u > v:
                u, v = v, u
            num[u] += 1
            num[v] += 1
            M[(u, v)] = M.get((u, v), 0) + 1
        
        for i in range(1, n + 1):
            f[i - 1] = num[i]
        f.sort()
        ans = [0] * q
        for i in range(1, n + 1):
            num1 = num[i]
            for j in range(0, q):
                cnt = queries[j]
                need = cnt - num1
                idx = bisect.bisect_right(f, need)
                ans[j] += n - idx
                if num1 + num1 > cnt:
                    ans[j] -= 1

        for key, val in M.items():
            u, v = key
            num1 = num[u] + num[v]
            for j in range(0, q):
                cnt = queries[j]
                if num1 > cnt and num1 - val <= cnt:
                    ans[j] -= 2
        
        for i in range(0, q):
            ans[i] //= 2
        
        return ans