K-th Smallest in Lexicographical Order — DSA

Problem#

Given integers n and k, return the k-th smallest number (1-indexed) in lexicographic order among [1, n].

Examples#

n = 13, k = 2 → 10.
n = 1, k = 1 → 1.

Constraints#

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

Hints#

Hint 1

Imagine the implicit digit trie. From prefix p, count how many integers in [1, n] have p as a prefix.

Hint 2

If that count fits in the remaining k, skip the whole subtree by adding 1 to the prefix; otherwise descend into prefix * 10.

Approach#

Walk down the digit trie. At node cur, compute steps = count of integers in [cur, n] whose decimal starts with cur (i.e., [cur, cur+1) ∪ [cur*10, (cur+1)*10) ∪ ...). If steps <= k, skip the whole subtree (cur++, k -= steps); else descend (cur *= 10, k--).

Solution#

class Solution {
public:
    int findKthNumber(int n, int k) {
        long long cur = 1;
        --k;
        while (k > 0) {
            long long steps = countSteps(n, cur, cur + 1);
            if (steps <= k) {
                k -= steps;
                ++cur;
            } else {
                --k;
                cur *= 10;
            }
        }
        return static_cast<int>(cur);
    }
private:
    long long countSteps(long long n, long long p1, long long p2) {
        long long steps = 0;
        while (p1 <= n) {
            steps += min(n + 1, p2) - p1;
            p1 *= 10;
            p2 *= 10;
        }
        return steps;
    }
};

func findKthNumber(n int, k int) int {
    cur := 1
    k--
    for k > 0 {
        steps := countSteps(n, cur, cur+1)
        if steps <= k {
            k -= steps
            cur++
        } else {
            k--
            cur *= 10
        }
    }
    return cur
}

func countSteps(n, p1, p2 int) int {
    steps := 0
    for p1 <= n {
        upper := n + 1
        if p2 < upper {
            upper = p2
        }
        steps += upper - p1
        p1 *= 10
        p2 *= 10
    }
    return steps
}

class Solution:
    def findKthNumber(self, n: int, k: int) -> int:
        cur = 1
        k -= 1
        while k > 0:
            steps = self._count_steps(n, cur, cur + 1)
            if steps <= k:
                k -= steps
                cur += 1
            else:
                k -= 1
                cur *= 10
        return cur

    def _count_steps(self, n: int, p1: int, p2: int) -> int:
        steps = 0
        while p1 <= n:
            steps += min(n + 1, p2) - p1
            p1 *= 10
            p2 *= 10
        return steps

function findKthNumber(n, k) {
    let cur = 1;
    k--;
    while (k > 0) {
        const steps = countSteps(n, cur, cur + 1);
        if (steps <= k) {
            k -= steps;
            cur++;
        } else {
            k--;
            cur *= 10;
        }
    }
    return cur;
}

function countSteps(n, p1, p2) {
    let steps = 0;
    while (p1 <= n) {
        steps += Math.min(n + 1, p2) - p1;
        p1 *= 10;
        p2 *= 10;
    }
    return steps;
}

class Solution {
    public int findKthNumber(int n, int k) {
        long cur = 1;
        k--;
        while (k > 0) {
            long steps = countSteps(n, cur, cur + 1);
            if (steps <= k) {
                k -= (int) steps;
                cur++;
            } else {
                k--;
                cur *= 10;
            }
        }
        return (int) cur;
    }

    private long countSteps(long n, long p1, long p2) {
        long steps = 0;
        while (p1 <= n) {
            steps += Math.min(n + 1, p2) - p1;
            p1 *= 10;
            p2 *= 10;
        }
        return steps;
    }
}

function findKthNumber(n: number, k: number): number {
    let cur = 1;
    k--;
    while (k > 0) {
        const steps = countSteps(n, cur, cur + 1);
        if (steps <= k) {
            k -= steps;
            cur++;
        } else {
            k--;
            cur *= 10;
        }
    }
    return cur;
}

function countSteps(n: number, p1: number, p2: number): number {
    let steps = 0;
    while (p1 <= n) {
        steps += Math.min(n + 1, p2) - p1;
        p1 *= 10;
        p2 *= 10;
    }
    return steps;
}

Editorial#

The crux is countSteps: integers in subtree cur form intervals [cur, cur+1), [cur*10, (cur+1)*10), [cur*100, (cur+1)*100), ... Each level contributes min(n+1, p2) - p1 numbers, clipped by n. We descend by ten when k is small enough to land inside the subtree, otherwise we skip past it.

Complexity#

Time: O((log n)^2) — at most log n levels of descent, each calling countSteps which is log n.
Space: O(1).

Concept revision#

Lexicographical Numbers

Problem#

Examples#

Constraints#

Hints#

Approach#

Solution#

Editorial#

Complexity#

Concept revision#

Related#