440. K-th Smallest in Lexicographical Order

https://leetcode.com/problems/k-th-smallest-in-lexicographical-order/

把数字看作字符串,问[1, n]所有数字按字典序排序后第K小的。第K小会想到quick select/heap或二分法,但这些做法都超时。这道题的最优解巧妙利用了字典序的性质,把数字根据顺序看作一棵树(有点像trie),10~19是1的子节点,2是1的兄弟,100~199是10的子节点,如此往下到小于或等于n的数字。然后先序遍历这颗树,先遇到的节点排在前面,当子树已遍历完还没有到K则回溯到没有遍历过的兄弟节点,如此直到遍历到第K个。这个思路还可以继续优化,如果预先知道两个兄弟节点之间相隔多少节点且K大于这个间隔,说明结果不会在第一个节点的子树里,可以直接跳过第一个节点到第二个节点,省去了遍历第一个节点的子树的时间。计算间隔只要按层遍历,每层之间的数值之间差了10倍,遍历直到超过n为止。

解法参考自

/*
 * @lc app=leetcode id=440 lang=cpp
 *
 * [440] K-th Smallest in Lexicographical Order
 */

// @lc code=start
class Solution {
public:
    int findKthNumber(int n, int k) {
        const auto find_gap = [](long a, long b, long n) {
            // a = 1, b = 2, n = 198, gap = 1
            // a = 13, b = 20 gap = 10 + 1
            // a = 100, b = 200, gap = 199 - 100 + 11 = 120
            long gap = 0;
            while (a <= n) {
                gap += min(n + 1, b) - a;
                a *= 10;
                b *= 10;
            }
            return gap;
        };

        long cur = 1;
        while (k > 1) {
            long gap = find_gap(cur, cur + 1, n);
            if (gap <= k - 1) {
                k -= gap;
                ++cur;
            } else {
                cur *= 10;
                --k;
            }
        }
        return (int)cur;
    }
};
// @lc code=end

Last updated