Prefix Tree

440. K-th Smallest in Lexicographical Order

Role of countSteps

• countSteps(n, prefix1, prefix2) counts numbers starting with prefix1 up to but not including prefix2, bounded by n.

• It sums counts at every level of the prefix tree by multiplying prefix1 and prefix2 by 10 each time to move to the next level (deeper digits).

Summary:

• The algorithm efficiently skips entire subtrees of the prefix tree based on step counts.

• Moves horizontally (curr++) when the entire subtree is before k-th number.

• Moves vertically (curr *= 10) to dive deeper into the prefix tree when the k-th number lies in that subtree.

public int findKthNumber(int n, int k) {
    int curr = 1;
    k--;

    while (k > 0) {
        int step = countSteps(n, curr, curr + 1);
        if (step <= k) {
            curr++;
            k -= step;
        } else {
            curr *= 10;
            k--;
        }
    }
    return curr;
}

// To count how many numbers exist between prefix1 and prefix2
private int countSteps(int n, long prefix1, long prefix2) {
    int steps = 0;
    while (prefix1 <= n) {
        steps += Math.min(prefix2, n + 1) - prefix1;
        prefix1 *= 10;
        prefix2 *= 10;
    }
    return steps;
}