[LC]2343.裁剪数字后查询第k小的数字

题目

给你一个下标从 0 开始的字符串数组 nums ，其中每个字符串 长度相等 且只包含数字。

再给你一个下标从 0 开始的二维整数数组 queries ，其中 queries[i] = [ki, trimi] 。对于每个 queries[i] ，你需要：

• 将 nums 中每个数字 裁剪 到剩下 最右边 trimi 个数位。
• 在裁剪过后的数字中，找到 nums 中第 ki 小数字对应的 下标 。如果两个裁剪后数字一样大，那么下标 更小 的数字视为更小的数字。
• 将 nums 中每个数字恢复到原本字符串。

请你返回一个长度与 queries 相等的数组 answer，其中 answer[i]是第 i 次查询的结果。

提示：

• 裁剪到剩下最右边 x 个数位的意思是不断删除最左边的数位，直到剩下 x 个数位。
• nums 中的字符串可能会有前导 0 。

分析

大顶堆

依次处理每个query，处理每个query的时候用大顶堆的方式找到第k小的数字。

重点

在自定义比较操作的时候，不要使用 substr，substr 会进行拷贝操作，开销较高，应该使用 std::string::compare 对子串进行比较。

代码

class Solution {
public:
    vector<int> smallestTrimmedNumbers(vector<string>& nums, vector<vector<int>>& queries) {
        vector<int> ans;
        for (const auto & query : queries)
        {
            ans.push_back(solveQuery(nums, query[1], query[0]));
        }
        return ans;
    }
    int solveQuery(vector<string>& nums, int trim, int k)
    {
        int n = nums[0].size();
        auto cmp = [trim, n, &nums](const int lhs, const int rhs)
        {
            int res = nums[lhs].compare(n - trim, trim, nums[rhs], n - trim);
            return res < 0 || (res == 0 && lhs < rhs);
        };
        priority_queue<int, vector<int>, decltype(cmp)> pq(cmp);
        for (int i = 0; i < nums.size(); ++i)
        {
            if (pq.size() < k) 
            {
                pq.push(i);
                continue;
            }
            if (cmp(i, pq.top()))
            {
                pq.pop();
                pq.push(i);
            }
        }
        return pq.top();
    }
};

时间复杂度

假设总共有$m$个query，数字数组大小为$n$，字符串的长度为$s$.

一个query的时间复杂度为$O(ntrim_ilogk_i).$

所以总的时间复杂度为$O(n\Sigma_{i=1}^mtrim_ilogk_i)$

空间复杂度：$O(k).$