440. Kth Smallest in Lexicographical Order

Problem Link

Tutorial

Imagine this problem as a problem of walking on a tree, the root has children 1~9, 1 has children 10~19 and so on. All the nodes are less than n.

The pre-order traversal of the tree is the lexicographcal order of all the numbers. So now we stand at node 1 and travel (k-1) times. Then the node we standing on is the answer. Everytime we will move to the siblings on the right handside, if we have enough moves to travel all the nodes in the current subtree. Otherwise we go to the first child and do the same thing. Repeat like this until we use up all the moves.

We can easily see it on the tree, everytime we travel to the right, ans should be updated by ans + 1. Everytime we travel to the first child, the ans should be updated to ans * 10.

Now let's see how to count how many nodes in the current subtree. Just count the nodes with each depth and add them up. It should not be hard since each layer is a continuous sequence of integers.

Solution

class Solution(object):
    def findKthNumber(self, n, k):
    """
    :type n: int
    :type k: int
    :rtype: int
    """
    ret = 1
    k -= 1
    while k > 0:
        a = ret
        b = ret + 1
        temp_sum = 0
        while a <= n:
            temp_sum += min(b - 1, n) - (a - 1)
            b *= 10
            a *= 10
        if k >= temp_sum:
            ret += 1
            k -= temp_sum
        else:
            ret *= 10
            k -= 1
    return ret