以TopK为例,分析如何逐步优化算法的时间复杂度

2,349 阅读7分钟

分享一下怎么逐步优化一个算法的时间复杂度,以一个比较简单和常见的TopK 问题入手.
LeetCode215举例

在未排序的数组中找到第 k 个最大的元素。请注意,你需要找的是数组排序后的第 k 个最大的元素,而不是第 k 个不同的元素。
示例 1:
输入: [3,2,1,5,6,4] 和 k = 2
输出: 5
示例 2:
输入: [3,2,3,1,2,4,5,5,6] 和 k = 4
输出: 4
说明:

你可以假设 k 总是有效的,且 1 ≤ k ≤ 数组的长度。

使用排序的解法

要查找一个数组中第K大的元素,最直接的思想就是将这个数组排序,然后取第K大的元素就可以.


    func findKthLargest(_ nums: [Int], _ k: Int) -> Int {
       return nums.sorted()[nums.count-k]
    }

这种解法是对整个数组进行了排序.它的时间复杂度是O(n*lg(n)),其实我们的需求是取出第K大的元素,但是却将整个数组都进行了排序,对一些本来不需要排序的元素做了排序,那我们按照这个思路,放弃对整个数组排序,只进行局部排序,选出前K个最大的元素.

我们使用冒泡法对前K大个元素进行排序
例如

 假设 arr = [3,6,9,7,8,1,2,3,5,6] ,k = 4    

     func findKthLargest(_ nums: [Int], _ k: Int) -> Int {
        var nums = nums
        let count = nums.count
        for i in 0..<k  {
            for j in 0..<count-i-1 {
                if nums[j] > nums[j+1] {
                    (nums[j] ,nums[j+1]) = (nums[j+1] ,nums[j])
                }
            }
        }
        return nums[count-k]
    }

那么第K大的数就是6.
这种写法的时间复杂度是O(K*N),时间复杂度的大小取决于K的大小.如果K太大这种写法的复杂度还是相当高.因为是找到第K大的元素.不一定要对前K个元素进行排序,我们尝试不使用排序的方法.

使用堆

获取一个数组中最大或者是最小的元素让人很容易联想起堆.为了更好的理解接下来的解法我们简单复习一下堆的几个最主要的概念.

堆的概念

堆就是用数组实现的二叉树,所有它没有使用父指针或者子指针。堆根据“堆属性”来排序,“堆属性”决定了树中节点的位置。

堆属性

堆分为两种:最大堆和最小堆,两者的差别在于节点的排序方式。

在最大堆中,父节点的值比每一个子节点的值都要大。在最小堆中,父节点的值比每一个子节点的值都要小。这就是所谓的“堆属性”,并且这个属性对堆中的每一个节点都成立。
最大堆示例图

最小堆示例图

堆在数组中的储存方式

父节点和子节点的数组索引公式:

parent(i) = (i-1)/2
left(i) = 2i + 1
right(i) = 2i + 2

示例图

堆的高度

具有n个节点的堆的高度为h = log2(n)

堆的基本操作

  • shiftUp():如果元素比其父元素更大(max-heap)或更小(min-heap),则需要与父元素交换, 这使元素向上移动。
  • shiftDown()。 如果元素比子元素小(max-heap)或更大(min-heap),这个操作使元素向下移动,也称为“堆化(heapify)”。
  • insert(value):将新元素添加到堆的末尾,然后使用shiftUp()来修复堆。
  • remove():删除并返回最大值(max-heap)或最小值(min-heap)。为了填充元素删除后留下的位置,让最后一个元素移动到根位置,然后使用shiftDown()修复堆。 (有时称为“提取最小值”或“提取最大值”。)
  • removeAtIndex(index):类似remove(),不仅可以删除根节点,也可以从堆中删除任何节点。如果新元素与其子元素不规整,则调用shiftDown();如果元素与其父元素不规整,则调用shiftUp()。
  • replace(index, value):为节点分配一个较小(min-heap)或较大(max-heap)的值。因为这会使堆属性失效,所以它使用shiftUp()来修复。 (也称为“减少键”和“增加键”。)

解题思路:

  1. 当我们使用最小堆的时候.我们选取提供数组中前K位元素,将其生成为一个最小堆.那堆顶的元素就是前K个元素中第K大的元素.

然后将数组中剩余的逐个与堆顶进行比较.如果大于堆顶就与其进行交换,然后shiftdown.那堆顶就是第K大元素.

Swift代码的实现
func findKthLargest(_ nums: [Int], _ k: Int) -> Int {
    
    var array :Array<Int>!
    var heap :Heap<Int>!
    if nums.count > k {
        array = Array.init(nums[0..<k])
        heap = Heap.init(array: array, sort: <)
        for item in nums[k...nums.count-1] {
            if item > heap.peek()! {
                heap.replace(index: 0, value: item)
            }
        }
    }else{
        array = nums
        heap = Heap.init(array: array, sort: <)
    }
    return heap.peek()!
}

时间复杂度分析:这是一个由数组生成的堆,需要将整个数组遍历一遍,假设每个元素运气很差,每次都入堆调整,调整时间复杂度为堆的高度,即lg(N),故整体时间复杂度是O(N*lgN). 2. 使用最大堆的解法,用数组生成一个最大堆,然后将堆顶移除k-1次,那堆顶就是第K大的元素

Swift代码的实现
func findKthLargest(_ nums: [Int], _ k: Int) -> Int {
    
    var heap = Heap.init(array: nums, sort: >)
    for _ in 0..<k-1 {
        heap.remove()
    }
    return heap.peek()!
}

用最大堆的解法,思路相对比较简单,但是复杂度要高一些.O(NlgN + KlgN)

使用二分搜索和快速排序的思想达到O(n)的解决方案

解题思路: 首先在数组内随机选择一个基准,然后围绕该基准对数组进行分区,大于等于基准的数放在基准值的左边,小于基准值的值放在基准的右边.然后返回基准值的索引,如果索引等于K,则基准值就是我们寻找的第K大的数.如果K小于基准值的索引则说明,第K大的值大于基准值,我们只需要在基准值的左侧重复这一过程,直到找到恰好位于第K位置的基准.如果k大于基准值则说明第K大的数小于基准值,我们只需要在基准值的右侧重复这一过程直到找到恰好位于第K位置的基准.
举例:

[1,-2,-4,-7,3,4,11,23,14,27]

假设K=4,那我们需要选择第4大的元素.

  1. 首先第一步我们随机从数组中选择一个数,假如是是14.
  2. 我们对数组进行分区(不是排序只需要将数组遍历一遍就行)大于等于14的元素放在14的左边,小于14的放在14的右边.
[23, 27, 14, -7, 3, 4, 11, 1, -2, -4]
 <---- larger  smaller  -->

所有大于等于14的值都在左侧,所有小于14的值都在右侧,基准值14的索引是2(也就是第三大),所以第四大的元素一定在14的右侧,我们现在继续在右侧寻找.我们可以忽略掉数组的剩余部分

[x, x, x, -7, 3, 4, 11, 1, -2, -4]

我们再次随机从数组的剩余部分选择一个基准数11,然后对数组进行分区.

[x, x, x, 11,-7, 3, 4, 1, -2, -4]

分区得到如下数组,11的索引是3,是第四大的元素,正好是我们要找的元素.
Swift代码的实现

func randomizedPartition(_ a: inout [Int], _ low: Int, _ high: Int) -> Int {
    let pivotIndex = Int.random(in: low...high)
    a.swapAt(pivotIndex, high)
    let pivot = a[high]
    var store_index = low
    for j in low..<high {
        if a[j] <= pivot {
            a.swapAt(store_index, j)
            store_index += 1
        }
    }
    a.swapAt(store_index, high)
    return store_index
}


func randomizedSelect<T: Comparable>(_ a: inout [T], _ low: Int, _ high: Int, _ k: Int) -> T {
  if low < high {
    let p = randomizedPartition(&a, low, high)
    if k == p {
      return a[p]
    } else if k < p {
      return randomizedSelect(&a, low, p - 1, k)
    } else {
      return randomizedSelect(&a, p + 1, high, k)
    }
  } else {
    return a[low]
  }
}

复杂度分析:

  • 时间复杂度平均复杂度是O(n),最坏是最坏情况 O(N2).
  • 空间复杂度是O(1).

总结TopK虽然是一个比较简单的算法问题但是他逐步优化复杂度的思想非常值得我们借鉴.