裁剪数字后查询第 K 小的数字

标签: 数组 字符串 分治 快速选择 基数排序 排序 堆（优先队列）

难度: Medium

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

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

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

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

提示：

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

示例 1：

输入：nums = ["102","473","251","814"], queries = [[1,1],[2,3],[4,2],[1,2]]
输出：[2,2,1,0]
解释：
1. 裁剪到只剩 1 个数位后，nums = ["2","3","1","4"] 。最小的数字是 1 ，下标为 2 。
2. 裁剪到剩 3 个数位后，nums 没有变化。第 2 小的数字是 251 ，下标为 2 。
3. 裁剪到剩 2 个数位后，nums = ["02","73","51","14"] 。第 4 小的数字是 73 ，下标为 1 。
4. 裁剪到剩 2 个数位后，最小数字是 2 ，下标为 0 。
   注意，裁剪后数字 "02" 值为 2 。

示例 2：

输入：nums = ["24","37","96","04"], queries = [[2,1],[2,2]]
输出：[3,0]
解释：
1. 裁剪到剩 1 个数位，nums = ["4","7","6","4"] 。第 2 小的数字是 4 ，下标为 3 。
   有两个 4 ，下标为 0 的 4 视为小于下标为 3 的 4 。
2. 裁剪到剩 2 个数位，nums 不变。第二小的数字是 24 ，下标为 0 。

提示：

  • 1 <= nums.length <= 100
  • 1 <= nums[i].length <= 100
  • nums[i] 只包含数字。
  • 所有 nums[i].length 的长度 相同 。
  • 1 <= queries.length <= 100
  • queries[i].length == 2
  • 1 <= ki <= nums.length
  • 1 <= trimi <= nums[0].length

进阶：你能使用 基数排序算法 解决此问题吗？这种解法的复杂度又是多少？

Submission

运行时间: 200 ms

内存: 16.4 MB

class Solution:
    def smallestTrimmedNumbers(self, nums: List[str], queries: List[List[int]]) -> List[int]:
        n = len(nums)
        @cache
        def get_rank(cut):
            # 返回在cut处切割后的大小与下标排列列表
            return sorted(range(n), key = lambda i:nums[i][-cut:])
        
        return [get_rank(trim)[k-1] for k, trim in queries]

Explain

本题解采用排序和缓存策略处理查询。首先，定义一个辅助函数 get_rank，该函数接收一个参数 cut，表示需要保留的数位数量。对于每个数字，get_rank 会截取最右边的 cut 个数字，并对这些裁剪后的数字进行排序（如果数字相同，则按原来的索引顺序排序）。排序的结果是索引列表，表示裁剪后数字从小到大的索引排序。使用 @cache 装饰器缓存已计算的结果，避免重复计算相同的裁剪操作。对于每个查询，直接从缓存中获取排序结果，并返回第 k 小元素的索引。

时间复杂度: O(m * n log n)

空间复杂度: O(m * n)

class Solution:
    def smallestTrimmedNumbers(self, nums: List[str], queries: List[List[int]]) -> List[int]:
        n = len(nums)  # 数组长度
        @cache  # 使用缓存装饰器以避免重复计算
        def get_rank(cut):
            # 返回在cut处切割后的大小与下标排列列表
            return sorted(range(n), key=lambda i: nums[i][-cut:])  # 通过对裁剪后的数字进行排序和获取索引
        
        return [get_rank(trim)[k-1] for k, trim in queries]  # 对每个查询，使用缓存结果进行快速查询

Explore

在裁剪数字时，如果数字的长度小于指定的裁剪长度`trimi`，则使用整个数字进行操作。这是因为我们无法从一个较短的数字中裁剪出比其本身长度更多的位数。因此，整个数字被视为裁剪后的结果。这种处理方式简单且符合逻辑，允许在不同长度数字间进行公平比较。

在本题解中，数字字符串是按照字典序进行排序的。因此，如果裁剪后的数字中包含前导0，这些0会影响排序结果。例如，'001'会在'010'之前，因为字典序中'0'小于'1'。这确实可能改变原有数字在未裁剪前的相对顺序。但由于所有数字都是以相同的规则进行裁剪和比较，这种改变是合理的，反映了裁剪后数字的自然排序顺序。

在题解中，我们使用了Python的装饰器`@cache`来实现缓存机制。这个装饰器是从functools模块提供的，它可以自动缓存函数的返回结果。当函数使用相同的参数值被调用时，`@cache`将不会重新计算函数，而是直接返回之前存储的结果。这样，对于相同的`trimi`值，函数`get_rank`的计算结果会被缓存并在后续查询中重用，从而避免重复计算和提高效率。

在题解的排序逻辑中，我们使用了lambda函数作为排序的关键函数。这个lambda函数首先根据数字字符串的裁剪结果排序，如果存在相同的裁剪结果，则自动依据提供给排序函数的索引（即列表中的位置）进行排序，因为`range(n)`会生成一个自然的索引列表。这种排序方法保证了在裁剪结果相同的情况下，元素将按照它们在原始数组中的位置进行排序。通常，这类排序操作内部使用的是稳定排序算法，如归并排序，保持了相等元素的相对顺序。