Count Number of Bad Pairs LeetCode solution in Python

1
class Solution:
2
    def countBadPairs(self, nums: List[int]) -> int:
3
        nums_len = len(nums)
4
        count_dict = dict()
5
        for i in range(nums_len):
6
            nums[i] -= i
7
            if nums[i] not in count_dict:
8
                count_dict[nums[i]] = 0
9
            count_dict[nums[i]] += 1
10

11
        count = 0
12
        for key in count_dict:
13
            count += math.comb(count_dict[key], 2)
14
        return math.comb(nums_len, 2) - count