K-th Smallest in Lexicographical Order

Tested Python solution for LeetCode 440 with 26 pytest cases. Generate a practice environment with lcpy.

LeetCode 440, Hard. Topics: Trie. View on LeetCode.

Generate this problem as a practice environment: tested reference solution, 26 parametrized pytest cases, and a playground notebook:

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

​

Problem

Given two integers n and k, return the k^th lexicographically smallest integer in the range [1, n].

​

Examples

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

​

Constraints

• 1 <= k <= n <= 10^9

​

Solution

Reference implementation from solution.py on GitHub, full suite in test_solution.py:

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

​

Complexity

Time	Space
O(log(n)^2)	O(1)

​

Tags

NeetCode All.