LC 3272. 统计好整数的数目

https://leetcode.cn/problems/find-the-count-of-good-integers/description/

这题应该算是一个基本统计+排列组合

最开始第三步我尝试枚举,但是这个肯定不行。想象一下如果k=1的话,那么所有的数都是满足的。另外就是这里n, k取值范围都比较小,最后提交的话可以打表,但是也没啥意思。

我这里计算特征可能有点搓,用了一个set来自做去重。感觉这里如果可以按照某种顺序搜索的话,应该是可以省略这个set的。最后计算排列组合的时候需要注意起始位置不能是0.

class Solution:
    def countGoodIntegers(self, n: int, k: int) -> int:
        ans = 0
        opts = set()

        import functools
        @functools.cache
        def C(n, m):
            if m == 0: return 1
            if n == m: return 1
            return C(n - 1, m - 1) + C(n - 1, m)

        def check(cnt):
            ans = 1
            base = n
            if cnt[0] != 0:
                ans *= C(base - 1, cnt[0])
                base -= cnt[0]

            for c in cnt[1:]:
                if c != 0:
                    ans *= C(base, c)
                    base -= c
            return ans

        def dfs(i, now, cnt):
            nonlocal ans
            j = n - 1 - i
            if j < i:
                if now % k == 0:
                    cnt = tuple(cnt)
                    if cnt not in opts:
                        opts.add(cnt)
                        ans += check(cnt)
                return

            start = 1 if i == 0 else 0
            bi, bj = 10 ** i, 10 ** j
            for d in range(start, 10):
                if i != j:
                    d2 = now + d * bi + d * bj
                    cnt[d] += 2
                    dfs(i + 1, d2, cnt)
                    cnt[d] -= 2
                else:
                    d2 = now + d * bi
                    cnt[d] += 1
                    dfs(i + 1, d2, cnt)
                    cnt[d] -= 1
            return

        dfs(0, 0, [0] * 10)
        return ans