K-th Smallest in Lexicographical Order

hard trie math

Problem

Given integers n and k, return the k-th lexicographically smallest integer in [1, n]. Treat numbers as strings: "1" < "10" < "100" < "2".

Inputn = 13, k = 2

Output10

Lex order: 1, 10, 11, 12, 13, 2, 3, …; the 2nd is 10.

def find_kth_number(n, k):
    def steps(prefix):
        # numbers in [prefix, prefix+1) prefix-tree, capped by n
        first, last, cnt = prefix, prefix, 0
        while first <= n:
            cnt += min(n, last) - first + 1
            first *= 10
            last = last * 10 + 9
        return cnt

    cur, k = 1, k - 1
    while k > 0:
        s = steps(cur)
        if s <= k:
            k -= s
            cur += 1
        else:
            k -= 1
            cur *= 10
    return cur

function findKthNumber(n, k) {
  function steps(prefix) {
    let first = prefix, last = prefix, cnt = 0;
    while (first <= n) {
      cnt += Math.min(n, last) - first + 1;
      first *= 10;
      last = last * 10 + 9;
    }
    return cnt;
  }
  let cur = 1; k -= 1;
  while (k > 0) {
    const s = steps(cur);
    if (s <= k) { k -= s; cur += 1; }
    else { k -= 1; cur *= 10; }
  }
  return cur;
}

class Solution {
    long steps(long prefix, int n) {
        long first = prefix, last = prefix, cnt = 0;
        while (first <= n) {
            cnt += Math.min((long) n, last) - first + 1;
            first *= 10;
            last = last * 10 + 9;
        }
        return cnt;
    }
    public int findKthNumber(int n, int k) {
        long cur = 1; k -= 1;
        while (k > 0) {
            long s = steps(cur, n);
            if (s <= k) { k -= s; cur += 1; }
            else { k -= 1; cur *= 10; }
        }
        return (int) cur;
    }
}

long long stepsHelper(long long prefix, int n) {
    long long first = prefix, last = prefix, cnt = 0;
    while (first <= n) {
        cnt += min((long long) n, last) - first + 1;
        first *= 10;
        last = last * 10 + 9;
    }
    return cnt;
}
int findKthNumber(int n, int k) {
    long long cur = 1; k -= 1;
    while (k > 0) {
        long long s = stepsHelper(cur, n);
        if (s <= k) { k -= s; cur += 1; }
        else { k -= 1; cur *= 10; }
    }
    return (int) cur;
}

Explanation

Listing all numbers in lexicographical order and counting to k would be far too slow for large n. Instead, we imagine the numbers as a 10-way tree (trie) where each prefix like 1 branches into 10, 11, ..., and walk it cleverly while counting how many numbers each branch holds.

The helper steps(prefix) counts how many numbers in [1, n] start with that prefix. It does this level by level: at each depth the prefix covers the range [first, last], and we add min(n, last) - first + 1, then go one digit deeper by multiplying by 10. This is the subtree size.

Starting at prefix 1, we repeatedly decide: if the current subtree has s numbers and that is small enough (s <= k), we can skip the whole subtree, subtract s from k, and move to the next sibling (cur += 1). Otherwise the answer is inside this subtree, so we descend by doing cur *= 10 and use up one count.

We keep going until k reaches zero, at which point cur is exactly the k-th number.

Example: n = 13, k = 2. Lex order is 1, 10, 11, 12, 13, 2, 3, ... Starting at 1, the subtree under 1 is large, so we descend to 10, and that turns out to be the 2nd number — the answer.

Time: O(log² n) Space: O(1)