Python - Tìm phần tử nhỏ nhất thứ k trong thời gian tuyến tính

Cho danh sách nums và k (0-based). Tìm phần tử nhỏ nhất thứ k (kth smallest) trong O(n) trung bình.

Ví dụ

Input: nums = [6, 4, 9, 3, 1], k = 2
Output: 4
Sắp xếp: [1, 3, 4, 6, 9], phần tử index 2 là 4.

Thuật toán: Quickselect / Heap size k+1

Dùng Max Heap kích thước k+1:

1. Đẩy k+1 phần tử đầu vào max heap
2. Với mỗi phần tử còn lại: nếu nhỏ hơn đỉnh heap, pop đỉnh và push phần tử mới
3. Cuối cùng, đỉnh heap (lớn nhất trong k+1 phần tử nhỏ nhất) là kết quả

Mã Python

from heapq import heappop, heappush

def solve(nums, k):
    maxHeap = []
    
    # Khởi tạo heap với k+1 phần tử đầu (lưu âm để thành max-heap)
    for i in range(k + 1):
        heappush(maxHeap, -nums[i])
    
    # Xử lý phần còn lại
    for i in range(k + 1, len(nums)):
        if nums[i] < -maxHeap[0]:  # Nhỏ hơn max hiện tại?
            heappop(maxHeap)
            heappush(maxHeap, -nums[i])
    
    return -maxHeap[0]

# Test
nums = [6, 4, 9, 3, 1]
k = 2
print(solve(nums, k))  # Output: 4

Phân tích độ phức tạp

• Thời gian: O(n log k) - n phần tử, mỗi push/pop O(log k)
• Không gian: O(k) - heap size k+1

Lưu ý: Độ phức tạp O(n log k) không phải O(n) tuyệt đối. Để đạt O(n) trung bình, dùng Quickselect (phân hoạch như QuickSort). Heap approach O(n log k) phù hợp khi k nhỏ hoặc cần streaming.

Quickselect (O(n) trung bình)

import random

def quickselect(nums, k):
    if len(nums) == 1:
        return nums[0]
    pivot = random.choice(nums)
    left = [x for x in nums if x < pivot]
    mid = [x for x in nums if x == pivot]
    right = [x for x in nums if x > pivot]
    
    if k < len(left):
        return quickselect(left, k)
    elif k < len(left) + len(mid):
        return pivot
    else:
        return quickselect(right, k - len(left) - len(mid))

Bài toán gốc: LeetCode 215 - Kth Largest Element in Array (tương tự, kth largest = (n-k)th smallest).