Next Greater Numerically Balanced Number LeetCode solution in Python

1
class Solution:
2
    def nextBeautifulNumber(self, n: int) -> int:
3
        n_digits = len(str(n))
4

5
        next_max = {
6
            1: [1],
7
            2: [22],
8
            3: [122, 333],
9
            4: [1333, 4444],
10
            5: [14444, 22333, 55555],
11
            6: [122333, 224444, 666666, 155555],
12
            7: [1224444, 2255555, 3334444, 1666666, 7777777],
13
        }
14

15
        if n >= int(str(n_digits) * n_digits):
16
            n_digits += 1
17
            return min(next_max[n_digits])
18

19
        ans = float("inf")
20
        for num in sorted(next_max[n_digits]):
21
            cands = set(permutations(str(num)))
22
            cands = sorted(map(lambda x: int("".join(x)), cands))
23

24
            loc = bisect.bisect(cands, n)
25
            if loc < len(cands):
26
                ans = min(ans, cands[loc])
27

28
        return ans