440.字典序的第K小数字

440.字典序的第K小数字

Thought

当我们按照例子上想找到13的第2个最小的数字的时候,我们首先回去找数字1: 1 10 11 12 13

然后找到10,也就是说我们首先需要知道:

1. 1 -> 1 10 11 12 13
2. 2 -> 2 20 21 22 23
3. …
4. N -> N N0 N1 N2 N3

我们可以构建出来树:

       root
      / | \
     1  2  3 ...
    /       \
  10        ...
  /  \
100 101 ...

$关键点在于$：如何跳过整个子树
我们从 curr = 1 开始，尝试不断跳到字典序下一个数字。如果当前子树（比如以 curr = 1 开头的所有数字：1,10,11,…19…）的节点数少于 k，说明第 k 小不在这个子树，可以整体跳过

通过这个函数来计算出, 以 prefix(1)为前缀的所有数字中, <=n的个数

def getCount(prefix: int, n: int) -> int:
    count = 0
    cur = prefix
    next = prefix + 1
    while cur <= n:
        count += min(n + 1, next) - cur
        cur *= 10
        next *= 10
    return count

getCount(1, 13) => len(1 10 11 12 13) => 5

if count <= k:
这里 count=5 > k=1，说明我们要找的第 k 小数字 在 curr=1 的子树里

curr *= 10 # curr = 10
k -= 1 # k = 0

现在我们走到了 curr = 10，k 也变成了 0，说明我们已经找到了第 k 小的数字。

Answer

class Solution:
    def findKthNumber(self, n: int, k: int) -> int:
        curr = 1
        k -= 1  # 已经在第一个位置了，所以提前减去 1

        while k > 0:
            count = self.getCount(curr, n)
            if count <= k:
                curr += 1  # 跳过整个子树
                k -= count
            else:
                curr *= 10  # 进入子树
                k -= 1
        return curr

    def getCount(self, prefix: int, n: int) -> int:
        count = 0
        cur = prefix
        next = prefix + 1
        while cur <= n:
            count += min(n + 1, next) - cur
            cur *= 10
            next *= 10
        return count