K-th Smallest in Lexicographical Order

lcpy gen -n 440   # by problem number
lcpy gen -s k_th_smallest_in_lexicographical_order   # by problem name

Input: n = 13, k = 2
Output: 10
Explanation: The lexicographical order is [1, 10, 11, 12, 13, 2, 3, 4, 5, 6, 7, 8, 9], so the second smallest number is 10.

Input: n = 1, k = 1
Output: 1

class Solution:
    # Time: O(log(n)^2)
    # Space: O(1)
    def find_kth_number(self, n: int, k: int) -> int:
        curr = 1
        k -= 1
        while k:
            steps = self._count_steps(n, curr)
            if steps <= k:
                # skip the whole subtree rooted at curr
                curr += 1
                k -= steps
            else:
                # descend into the subtree
                curr *= 10
                k -= 1
        return curr

    def _count_steps(self, n: int, prefix: int) -> int:
        # count numbers in [1, n] starting with `prefix`
        steps = 0
        first = prefix
        last = prefix
        while first <= n:
            steps += min(last, n) - first + 1
            first *= 10
            last = last * 10 + 9
        return steps