剑指offer_40_最小的 K 个数

Posted 2020-12-18 小片清风

最小的 k 个数

题目链接：https://leetcode-cn.com/problems/zui-xiao-de-kge-shu-lcof/

题目内容：输入整数数组 arr ，找出其中最小的 k 个数。例如，输入4、5、1、6、2、7、3、8这8个数字，则最小的4个数字是1、2、3、4。

示例 1：

输入：arr = [3,2,1], k = 2
输出：[1,2] 或者 [2,1]
示例 2：

输入：arr = [0,1,2,1], k = 1
输出：[0]

限制：

• 0 <= k <= arr.length <= 10000
• 0 <= arr[i] <= 10000

题目解析

题目的意思简单明了。获取最小的前 k 个数。

题目解法

方法一：python 内置函数排序

时间复杂度 O(nlogn) python 内置函数时间复杂度点此查看

代码：

class Solution:
    def getLeastNumbers(self, arr: List[int], k: int) -> List[int]:
        arr.sort()
        return arr[:k]

面试就别拿这个献丑了，知道python牛就得了

方法二：手动排序，取前 k 小。快排为宜

为什么用快排呢。就是在取前 k 小的时候，可以判断在第 k 小的值出现在左侧(中间值 - i + 1 > k)时，向左递归，出现在右侧(中间值pos - i + 1 < k), 向右递归。

时间复杂度最好是O(n), 最坏会退化到O(n^2)

代码：

class Solution:
    def getLeastNumbers(self, arr: List[int], k: int) -> List[int]:
        def partition(arr, l, r):
            #选定中值
            pivotvalue = arr[l]
            lmark = l + 1
            rmark = r
            done = False

            while not done:
                while lmark <= rmark and arr[lmark] <= pivotvalue:
                    lmark += 1
                while rmark >= lmark and arr[rmark] >= pivotvalue:
                    rmark -= 1
                if rmark < lmark:
                    done = True
                else:
                    arr[lmark], arr[rmark] = arr[rmark], arr[lmark]

            arr[l], arr[rmark] = arr[rmark], arr[l]
            return rmark
        
        def quicksort(arr, l, r, k):
            if l > r:
                return 
            pos = partition(arr, l, r)
            num = pos - l + 1
            if k == num:
                return
            if k < num:
                quicksort(arr, l, pos - 1, k)
            else:
                quicksort(arr, pos+1, r, k - num)

        if k == 0:
            return []
        quicksort(arr, 0, len(arr) - 1, k)
        return arr[:k]

方法三：使用大顶堆

求前 k 个最小用大顶堆。为什么呢？因为当我们根据题意维护一个 K 大小的大顶堆时，向堆中添加元素时，当这个被添加元素小于堆顶值(也就是当前堆最大值)时，大顶堆会抛出最大值，将这个值插入大顶堆。这样，我们就得到了比堆顶小的 k - 1 个最小值。这样逐个插入元素时，最后得到的堆里存的就是最小的 k 个数。

时间复杂度O(nlogk), 因为我们只需要维护 K 大小的堆。

代码：

class HeapList(object):
    """大顶推"""
    def __init__(self):
        self.heaplist = [0]
        self.size = 0

    def buildHeap(self, alist):
        i = len(alist) // 2
        self.size = len(alist)
        self.heaplist += alist[:]
        while i > 0:
            self.percDown(i)
            i -= 1

    def percUp(self, i):
        while i // 2 > 0:
            if self.heaplist[i] > self.heaplist[i // 2]:
                self.heaplist[i], self.heaplist[i // 2] = self.heaplist[i // 2], self.heaplist[i]
            i //= 2

    def insert(self, k):
        self.heaplist.append(k)
        self.size += 1
        self.percUp(self.size)

    def maxChild(self, i):
        if i * 2 + 1 > self.size:
            return i * 2
        else:
            if self.heaplist[i * 2] > self.heaplist[i * 2 + 1]:
                return i * 2
            else:
                return i * 2 + 1

    def percDown(self, i):
        while i * 2 <= self.size:
            mc = self.maxChild(i)
            if self.heaplist[i] < self.heaplist[mc]:
                self.heaplist[i], self.heaplist[mc] = self.heaplist[mc], self.heaplist[i]
            i = mc

    def delMax(self):
        retval = self.heaplist[1]
        self.heaplist[1] = self.heaplist[self.size]
        self.size -= 1
        self.heaplist.pop()
        self.percDown(1)
        return retval


# 采用大顶堆的方式，制作容量为 k 的大顶堆，向堆中添加元素时，比堆顶值小，就弹出堆顶，并将此元素添加进堆。这就保证，最后遍历完成后，
# 我们获得了比堆顶小的 k-1 个最小值
# 时间复杂度 O(nlogK)  因为只维护 K 大小的堆
class Solution:
    def getLeastNumbers(self, arr, k):
        if k == 0:
            return []
        heaplist = HeapList()
        heaplist.buildHeap(arr[:k])
        for i in arr[k: ]:
            if i < heaplist.heaplist[1]:
                heaplist.delMax()
                heaplist.insert(i)
        return heaplist.heaplist[1:]

附录

这里附一个 小顶堆的实现。

class Heap:
    """二叉堆的实现"""
    def __init__(self):
        self.heapList = [0]   # 默认一个 0 做占位，使得根节点的索引在 1 上
        self.currentSize = 0    # 最大节点的索引位置

    def perUp(self, i):
        """将小节点逐步上升"""
        while i // 2 > 0:
            if self.heapList[i] < self.heapList[i // 2]:
                self.heapList[i], self.heapList[i // 2] = self.heapList[i // 2], self.heapList[i]
            i = i // 2

    def insert(self, k):
        """插入节点"""
        self.heapList.append(k)
        self.currentSize += 1
        self.perUp(self.currentSize)

    def minChild(self, i):
        """获取左右两个子节点里较小的那个子节点的索引"""
        if i * 2 + 1 > self.currentSize:  # 右子节点超出节点数量
            return i * 2
        else:
            if self.heapList[i * 2] < self.heapList[i * 2 + 1]:
                return i * 2
            else:
                return i * 2 + 1

    def perDown(self, i):
        """将节点下沉到合适位置"""
        while (i * 2) <= self.currentSize:  # 说明有子节点
            mc = self.minChild(i)
            if self.heapList[i] > self.heapList[mc]:
                self.heapList[i], self.heapList[mc] = self.heapList[mc], self.heapList[i]
            i = mc

    def delMin(self):
        """删除小节点"""
        retval = self.heapList[1]  # 删除索引位置为 1 的节点
        self.heapList[1] = self.heapList[self.currentSize]
        self.heapList.pop()
        self.currentSize -= 1
        self.perDown(1)
        return retval

    def buildHeap(self, alist):
        i = len(alist) // 2
        self.currentSize = len(alist)
        self.heapList += alist[:]
        while i > 0:
            self.perDown(i)
            i -= 1