标签: 树状数组 线段树 数组 二分查找 分治 有序集合 归并排序
难度: Hard
标签: 树状数组 线段树 数组 二分查找 分治 有序集合 归并排序
难度: Hard
运行时间: 112 ms
内存: 23.0 MB
class Solution: def kBigIndices(self, nums: List[int], k: int) -> int: # at least k different idx1 idx1 < i and nums[idx1] < nums[i] # at least k different idx2 idx2 > i and nums[idx2] < nums[i] # use priority queue iterate from left to right, and iterate from right to left n = len(nums) if n <= 2 * k: return 0 count_big = 0 count = [False] * (n) pq = [] for index, val in enumerate(nums): if len(pq) == k and val > -pq[0]: count[index] = True heappush(pq, -val) #use max heap the biggest number on the top if len(pq) > k: heappop(pq) print(count) pq = [] reversed_nums = list(reversed(nums)) count = count[::-1] print(reversed_nums) for index, val in enumerate(reversed_nums): if len(pq) == k and val > -pq[0] and count[index]: count_big += 1 heappush(pq, -val) if len(pq) > k: heappop(pq) return count_big
题解的整体思路是使用两次遍历和两个优先队列(最大堆)来计算每个索引 i,左边至少有 k 个元素小于 nums[i] 且右边至少有 k 个元素小于 nums[i] 的索引数量。首先,从左到右遍历数组,对于每个元素,使用最大堆来维护它左边最大的 k 个元素。如果堆的大小达到了 k 并且当前元素大于堆顶元素,标记该索引。然后,从右到左遍历数组,使用另一个最大堆来维护每个元素右边最大的 k 个元素。如果堆的大小达到了 k,当前元素大于堆顶元素,并且该索引在前一次遍历中已被标记,则该索引符合条件,计数增加。最后返回符合条件的索引数量。
时间复杂度: O(n log k)
空间复杂度: O(n + k)
class Solution: def kBigIndices(self, nums: List[int], k: int) -> int: n = len(nums) if n <= 2 * k: return 0 count_big = 0 count = [False] * (n) pq = [] # 从左到右遍历数组,使用最大堆记录左侧最大的k个数 for index, val in enumerate(nums): if len(pq) == k and val > -pq[0]: count[index] = True heappush(pq, -val) if len(pq) > k: heappop(pq) pq = [] reversed_nums = list(reversed(nums)) count = count[::-1] # 从右到左遍历数组,使用最大堆记录右侧最大的k个数 for index, val in enumerate(reversed_nums): if len(pq) == k and val > -pq[0] and count[index]: count_big += 1 heappush(pq, -val) if len(pq) > k: heappop(pq) return count_big
在算法中,我们使用最大堆来维护每个元素左侧的最大k个元素。当堆的大小达到k时,堆顶元素是这k个元素中最小的,即是第k大的元素。如果当前元素大于这个堆顶元素,意味着当前元素至少比左侧的这k个元素都要大,从而确保存在至少k个元素小于当前元素。
选择使用优先队列(最大堆)主要是因为它能高效地维护一组元素中的最大值或最小值。在本题中,需要快速获取并更新每个元素左侧和右侧的最大k个元素。最大堆能够在对元素进行插入和删除操作时,保持常数时间内获取最大值(堆顶元素),这对于检查当前元素是否大于左侧或右侧的k个元素中的最小一个(即堆顶元素)是非常有效的。
首先从左到右遍历是为了确定每个索引位置的元素是否大于其左侧的k个元素。接着,从右到左遍历则是为了确定这些元素是否也大于其右侧的k个元素。只有当一个元素同时满足这两个条件时,它才被计入最终的统计结果中。这种双向遍历确保了每个位置的元素都被正确地评估了其左侧和右侧的条件。