Taro Logo

K-th Smallest in Lexicographical Order

Hard
Amazon logo
Amazon
Topics:
TreesGreedy Algorithms

Given two integers n and k, return the kth lexicographically smallest integer in the range [1, n]. For example:

  • If n = 13 and k = 2, the lexicographical order is [1, 10, 11, 12, 13, 2, 3, 4, 5, 6, 7, 8, 9], so the function should return 10.
  • If n = 1 and k = 1, the function should return 1.

Write an efficient algorithm to solve this problem, considering that 1 <= k <= n <= 10^9.

Solution


Naive Approach

A brute-force solution would involve generating all numbers from 1 to n, sorting them lexicographically, and then returning the k-th element. This approach is highly inefficient due to the sorting step, which would have a time complexity of O(n log n).

Code (Python)

def findKthNumber_naive(n: int, k: int) -> int:
    nums = list(map(str, range(1, n + 1)))
    nums.sort()
    return int(nums[k - 1])

Time Complexity

O(n log n) due to sorting.

Space Complexity

O(n) to store the numbers in a list.

Optimal Solution: Lexicographical Tree Traversal

A more efficient approach involves treating the numbers from 1 to n as nodes in a 10-ary tree (lexicographical tree). We can traverse this tree to find the k-th smallest number without explicitly generating and sorting all numbers.

  1. calSteps(n, n1, n2) Function: This helper function calculates how many steps it takes to go from node n1 to node n2 (inclusive) in the lexicographical tree. This essentially counts the number of nodes between n1 and n2.
  2. Main Loop: Start with curr = 1 (the root of the tree). In each iteration, calculate the steps from curr to curr + 1. If the number of steps is less than or equal to k, it means the k-th number is not in the subtree rooted at curr, so we move to the next sibling (curr += 1) and reduce k by the number of steps we skipped. Otherwise, the k-th number is in the subtree rooted at curr, so we move to the first child (curr *= 10) and reduce k by 1 (since we are visiting the current node).
  3. Termination: The loop continues until we find the k-th number.

Code (Python)

def findKthNumber(n: int, k: int) -> int:
    def calSteps(n: int, n1: int, n2: int) -> int:
        steps = 0
        while n1 <= n:
            steps += min(n + 1, n2) - n1
            n1 *= 10
            n2 *= 10
        return steps

    curr = 1
    k -= 1
    while k > 0:
        steps = calSteps(n, curr, curr + 1)
        if steps <= k:
            curr += 1
            k -= steps
        else:
            curr *= 10
            k -= 1
    return curr

Time Complexity

O(log n * log n). The outer while loop iterates at most log(n) times, as curr increases by at least a factor of 10 in each iteration when it multiplies by 10, and increases by one if not. calSteps takes O(log n) time.

Space Complexity

O(1). The algorithm uses only a constant amount of extra space.

Edge Cases

  • n = 1, k = 1: The algorithm should return 1.
  • n is a power of 10: The algorithm should correctly traverse the tree.
  • k = 1: The algorithm should return 1.