[剑指 Offer 第 2 版第 40 题] “最小的 K 个数”做题记录

第 40 题：最小的 K 个数

传送门：最小的k个数，牛客网 online judge 地址。

输入 n 个整数，找出其中最小的 k 个数。

注意：

• 数据保证 k 一定小于等于输入数组的长度;
• 输出数组内元素请按从小到大顺序排序;

样例：

输入：[1,2,3,4,5,6,7,8], k=4

输出：[1,2,3,4]

分析：最简单的思路就是排个序，然后取出前 k 个元素，不过时间复杂度是 O(n\\log n)，n 为数组的长度。

Python 代码：

  class Solution(object):
        def getLeastNumbers_Solution(self, input, k):
            size = len(input)
            if size == 0:
                return []
            if k == size:
                return sorted(input)
            return sorted(input)[:k]

其实可以用 O(n\\log k) 的时间复杂度找到最小的 k 个元素。有两种思路：

1、借助快速排序把数组一分为二的 partition 操作；

2、借助最大堆（需要把数组做一个转换，都变成相反数，最小的 k 个数，就是最大堆里最大的 k 个数）。

Python 代码1：partition，注意，这种方式，大于等于 pivot 的元素都被分在了右边

  class Solution(object):
        def getLeastNumbers_Solution(self, input, k):
            """
            :type input: list[int]
            :type k: int
            :rtype: list[int]
            """
            size = len(input)

            if size == 0:
                return []
            if k == size:
                return sorted(input)
            l = 0
            r = size - 1
            while l <= r:
                p = self.__partition(input, l, r)
                if p == k - 1:
                    return sorted(input[:p + 1])
                elif p > k - 1:
                    # 此时  k-1  p
                    r = p - 1
                else:
                    # 此时 p k-1
                    l = p + 1

        def __partition(self, input, left, right):
            # 只有一个数的时候，就没有必要 partition 了
            # 直接返回这个数的索引
            if left == right:
                return left
            pivot = input[left]

            j = left
            # [left + 1, j] 这个区间里的元素都严格小于 pivot
            for i in range(left + 1, right + 1):
                if input[i] < pivot:
                    j += 1
                    input[i], input[j] = input[j], input[i]
            input[left], input[j] = input[j], input[left]
            return j


    if __name__ == '__main__':
        input = [9, 14, 1, 16, 19, 13, 12]
        k = 4
        solution = Solution()
        result = solution.getLeastNumbers_Solution(input, k)
        print(result)

Python 代码2：最大堆

  class Solution(object):

        def getLeastNumbers_Solution(self, input, k):
            """
            :type input: list[int]
            :type k: int
            :rtype: list[int]
            """
            size = len(input)
            if size == 0:
                return []
            if k == size:
                return sorted(input)
            import heapq

            l = []
            for num in input[:k]:
                heapq.heappush(l, -num)
            for num in input[k:]:
                top = l[0]
                if top < -num:
                    heapq.heappushpop(l, -num)
            return sorted([-num for num in l])

当然，Python 中的 heapq 直接就有获取最小 k 个元素的方法。

Python 代码：

  class Solution(object):

        def getLeastNumbers_Solution(self, input, k):
            """
            :type input: list[int]
            :type k: int
            :rtype: list[int]
            """
            size = len(input)
            if size == 0:
                return []
            if k == size:
                return sorted(input)
            import heapq
            heapq.heapify(input)
            return sorted(heapq.nsmallest(k, input))

Python 代码：两路 partition 的写法：

  class Solution(object):
        def getLeastNumbers_Solution(self, input, k):
            """
            :type input: list[int]
            :type k: int
            :rtype: list[int]
            """

            size = len(input)

            if size == 0:
                return []
            if k == size:
                return sorted(input)
            l = 0
            r = size - 1
            while l <= r:
                p = self.__partition(input, l, r)
                if p == k - 1:
                    return sorted(input[:p + 1])
                elif p > k - 1:
                    # 此时  k-1  p
                    r = p - 1
                else:
                    # 此时 p k-1
                    l = p + 1

        def __partition(self, input, left, right):
            # 只有一个数的时候，就没有必要 partition 了
            # 直接返回这个数的索引
            if left == right:
                return left
            pivot = input[left]
            l = left + 1
            r = right
            while True:
                while l <= right and input[l] <= pivot:
                    l += 1
                while r > left and input[r] >= pivot:
                    r -= 1
                if l > r:
                    break
                input[l], input[r] = input[r], input[l]
                l += 1
                r -= 1
            input[left], input[r] = input[r], input[left]
            return r

C++ 代码：

作者：liweiwei1419

来源：https://liweiwei1419.github.io/sword-for-offer/