3544. Count Almost Equal Pairs II

class Solution:
    def countPairs(self, nums: List[int]) -> int:
        from collections import defaultdict
        
        def get_all_variants(num: int) -> set:
            """Generate all possible numbers from 0, 1, or 2 swaps"""
            # Pad to same length (max 7 digits for 10^7)
            s = str(num).zfill(7)
            n = len(s)
            
            variants = set()
            
            # Original (0 swaps)
            variants.add(s)
            
            # 1 swap
            for i in range(n):
                for j in range(i + 1, n):
                    lst = list(s)
                    lst[i], lst[j] = lst[j], lst[i]
                    variants.add(''.join(lst))
            
            # 2 swaps - apply swap to all 1-swap variants
            one_swap_variants = list(variants)
            for v in one_swap_variants:
                for i in range(n):
                    for j in range(i + 1, n):
                        lst = list(v)
                        lst[i], lst[j] = lst[j], lst[i]
                        variants.add(''.join(lst))
            
            return variants
        
        # For each number, compute all variants
        n = len(nums)
        
        # Store counts using a map
        count = 0
        seen = defaultdict(int)
        
        for num in nums:
            # Get the padded version
            s = str(num).zfill(7)
            
            # Count pairs: how many previous numbers can become equal to this one?
            # A previous number x can match current number y if:
            # - One of y's variants equals x's canonical form (padded string)
            # But actually we need: some variant of x equals some variant of y
            # This is equivalent to: x's canonical is in y's variants OR y's canonical is in x's variants
            # Actually more precisely: the intersection of their variant sets is non-empty
            
            # But that's expensive. Instead:
            # For this number, count how many previous numbers have a matching canonical in our variants
            
            variants = get_all_variants(num)
            
            # Count matches with previous numbers
            for v in variants:
                count += seen[v]
            
            # Add this number's canonical to seen
            seen[s] += 1
        
        return count