LeetCode #215 Kth Largest: sort, heap และ quickselect เลือกเมื่อไร

Kth Largest Element #215 ให้หาค่าลำดับที่ k เมื่อเรียงจากมากไปน้อย นับค่าซ้ำด้วย เช่น [3, 2, 3, 1, 2, 4, 5, 5, 6] กับ k=4 ตอบ 4 ไม่ใช่ค่าต่างกันลำดับสี่ โจทย์ถามเพียงค่าเดียว จึงมีทางเลือกดีกว่าการเรียงทั้งหมดในหลายสถานการณ์

ข้อจำกัดคือ 1 <= k <= n <= 100000 และค่าอยู่ระหว่าง -10000 ถึง 10000 ลองนึกภาพหาผู้เข้าเส้นชัยอันดับ k: จะเรียงทั้งสนาม, เก็บเฉพาะผู้ที่ติดอันดับ, หรือแบ่งสนามรอบจุด pivot ก็ได้

วิธี 1: sort สำเนา

sorted(nums, reverse=True)[k-1] ตรงไปตรงมา เวลา O(n log n) และมีสำเนา O(n) เพื่อไม่แก้ input ดีเมื่อ n เล็กหรืออยากได้ลำดับอื่นต่อ แต่ทำงานเกินจำเป็นถ้าขอเพียงหนึ่งค่า

วิธี 2: min-heap ขนาด k

เก็บ k ค่าที่ใหญ่ที่สุดใน heap โดยค่าบนสุดเป็นอันดับ k ปัจจุบัน เวลา O(n log k) พื้นที่ O(k) อ่าน input แบบ stream ได้และไม่แก้ input เหมาะเมื่อ k เล็กเทียบ n; ถ้า k ใกล้ n heap ใหญ่ขึ้นตาม

วิธี 3: quickselect แบบสามกลุ่ม

แบ่งค่าที่น้อยกว่า เท่ากับ และมากกว่า pivot แล้วโฟกัสเฉพาะช่วงที่มี index เป้าหมาย เวลาเฉลี่ย O(n) พื้นที่เพิ่ม O(1) เมื่อแก้ input ได้ แต่กรณี pivot แย่ซ้ำ ๆ อาจเป็น O(n²) วิธีนี้รองรับค่าซ้ำได้ด้วยกลุ่มเท่ากัน

import heapq
import unittest


def kth_sort(nums: list[int], k: int) -> int:
    return sorted(nums, reverse=True)[k - 1]


def kth_heap(nums: list[int], k: int) -> int:
    heap: list[int] = []
    for value in nums:
        if len(heap) < k:
            heapq.heappush(heap, value)
        elif value > heap[0]:
            heapq.heapreplace(heap, value)
    return heap[0]


def kth_quickselect_inplace(nums: list[int], k: int) -> int:
    target = len(nums) - k
    left, right = 0, len(nums) - 1
    while left <= right:
        pivot = nums[(left + right) // 2]
        low, i, high = left, left, right
        while i <= high:
            if nums[i] < pivot:
                nums[low], nums[i] = nums[i], nums[low]
                low += 1
                i += 1
            elif nums[i] > pivot:
                nums[high], nums[i] = nums[i], nums[high]
                high -= 1
            else:
                i += 1
        if target < low:
            right = low - 1
        elif target > high:
            left = high + 1
        else:
            return nums[target]
    raise AssertionError("valid k must return")


class KthLargestTests(unittest.TestCase):
    def test_all_methods_and_duplicates(self):
        cases = [([3, 2, 1, 5, 6, 4], 2, 5),
                 ([3, 2, 3, 1, 2, 4, 5, 5, 6], 4, 4),
                 ([2, 2, 2], 3, 2), ([-5], 1, -5),
                 ([5, 4, 3, 2, 1], 5, 1)]
        for values, k, expected in cases:
            self.assertEqual(kth_sort(values, k), expected)
            self.assertEqual(kth_heap(values, k), expected)
            self.assertEqual(kth_quickselect_inplace(values.copy(), k), expected)


if __name__ == "__main__":
    unittest.main()

Decision table

สถานการณ์	Sort copy	Heap k	Quickselect in-place
Input เล็ก/เน้นอ่านง่าย	ดีที่สุด	พอได้	โค้ดยาว
Memory จำกัด	O(n)	O(k)	O(1) เพิ่ม
Stream ที่มาเป็นช่วง	ต้องรอครบ	รับทีละค่าได้	ต้องมี array
Input เรียงแล้ว	อ่านตำแหน่งได้	ยังเดินทั้งหมด	ไม่จำเป็น
ห้ามแก้ input	ทำสำเนา	ไม่แก้	ต้องทำสำเนา O(n)
ต้องการ worst-case ที่คาดได้	O(n log n)	O(n log k)	อาจ O(n²)

สำหรับข้อจำกัดโจทย์นี้ heap เป็น default ที่ไว้ใจง่าย เมื่อขอค่าเดียวและไม่อยากแก้ input; ถ้า k ใหญ่มากและยอมแก้ input quickselect อาจเร็วกว่าโดยเฉลี่ย แต่ต้อง benchmark และยอมรับ worst case การเลือก sort ก็สมเหตุสมผลถ้า n เล็กหรือทีมให้ความสำคัญกับ code review มากกว่า micro-optimization

Edge cases และ Recap

• ค่าเท่ากันนับเป็นหลายตำแหน่ง ไม่ deduplicate
• k=1 คือค่าสูงสุด, k=n คือค่าต่ำสุด
• Heap ต้องเก็บไม่เกิน k ค่า; quickselect ในตัวอย่างแก้ list ต้นทาง
• เลือกวิธีตาม n, k, memory, stream และสัญญาการแก้ input