劍指offer：最小的k個數

題目描述
輸入n個整數，找出其中最小的K個數。例如輸入4,5,1,6,2,7,3,8這8個數字，則最小的4個數字是1,2,3,4,。

# -*- coding: utf-8 -*-
# @Time         : 2019-07-08 21:09
# @Author       : Jayce Wong
# @ProjectName  : job
# @FileName     : getLeastNumbers.py
# @Blog         : https://blog.51cto.com/jayce1111
# @Github       : https://github.com/SysuJayce

class Solution:
    """
    從數組中尋找最小的k個數字，通常來講最直接的作法就是先將這個數組排序，而後取最小的k個數字便可。
    可是這樣作的時間複雜度爲O(nlogn)

    解法1：
    借鑑快排中partition()的作法，由於partition()每次能夠肯定一個下標的正確元素，並保證其左右與其
    大小關係有序。因此只要咱們經過partition()肯定了下標爲k-1的元素，那麼其左邊都是小於該元素的。
    時間複雜度爲O(n)

    解法2：
    能夠維護一個容量爲k的容器，而後遍歷整個數組，若是遇到比容器中最大值要小的元素，那麼就將這個元素
    替換容器中的最大值。時間複雜度爲O(nlogk)
    """
    def GetLeastNumbers_Solution1(self, tinput, k):
        if k <= 0 or k > len(tinput):
            return []

        ans = tinput[:k]
        for i in range(k, len(tinput)):
            if tinput[i] < max(ans):
                ans.remove(max(ans))
                ans.append(tinput[i])

        return sorted(ans)

    def GetLeastNumbers_Solution2(self, tinput, k):
        def partition(begin, end):
            pos = begin
            for i in range(begin, end):
                if tinput[i] < tinput[end]:
                    tinput[i], tinput[pos] = tinput[pos], tinput[i]
                    pos += 1
            tinput[end], tinput[pos] = tinput[pos], tinput[end]
            return pos

        if k <= 0 or k > len(tinput):
            return []

        start, stop = 0, len(tinput) - 1
        idx = partition(start, stop)
        while idx != k - 1:
            if idx > k:
                stop = idx - 1
            else:
                start = idx + 1
            idx = partition(start, stop)

        return sorted(tinput[: k])

def main():
    nums = [4, 5, 1, 6, 2, 7, 3, 8]
    k = 4
    solution = Solution()
    print(solution.GetLeastNumbers_Solution2(nums, k))

if __name__ == '__main__':
    main()