On this page

ข้อ 49 · LC215 Kth Largest Element in an Array (ตัวมากอันดับ k) 🟡

👋 อ่านฟรีทั้งหมดบน Aph's Blog — เนื้อหาภาษาไทย ทำตามทีละหน้าใน sidebar ได้เลย หากมีข้อเสนอแนะหรืออยากให้เพิ่มหัวข้อไหน บอกได้เสมอ

หาค่าที่มากเป็นอันดับที่ k โดยไม่ต้อง sort (เรียง) ทั้ง array ใช้ min-heap ขนาด k เป็นกรอบเก็บ top-k

โจทย์ (LC215): กำหนด array จำนวนเต็ม nums และเลขจำนวนเต็ม k ให้ return ค่าที่มากเป็นอันดับที่ k ของ array เมื่อเรียงลำดับจากมากไปน้อย (นับตามลำดับการเรียง ไม่ใช่อันดับที่ k ของค่าที่ไม่ซ้ำกัน)

Example 1

Input:: nums = [3, 2, 1, 5, 6, 4], k = 2
Output:: 5
Explanation:: เรียงจากมากไปน้อยได้ [6, 5, 4, 3, 2, 1] อันดับ 1 คือ 6 อันดับ 2 คือ 5

Example 2

Input:: nums = [3, 2, 3, 1, 2, 4, 5, 5, 6], k = 4
Output:: 4
Explanation:: เรียงจากมากไปน้อยได้ [6, 5, 5, 4, 3, 3, 2, 2, 1] อันดับ 4 (นับค่าซ้ำด้วย) คือ 4

Constraints (ข้อจำกัด)

1 <= k <= nums.length <= 10^5
-10^4 <= nums[i] <= 10^4

แนวทาง — ต้องใช้อะไร & คิดยังไง

โครงสร้างที่ใช้: min-heap (จาก heapq) เราต้องการ maintain (รักษา) k ตัวที่มากที่สุดไว้ แล้วในบรรดา k ตัวนั้น ตัว minimum ก็คือคำตอบ (ตัวมากอันดับ k) พอดี

คิดแบบง่าย/ช้าก่อน: วิธี naive คือ sort ทั้ง array แล้วหยิบ element (สมาชิก) ที่ตำแหน่ง k จากท้าย ซึ่งเป็น O(n log n) แต่ถ้า k เล็กมากเทียบกับ n เราไม่จำเป็นต้อง sort ทุกตัว แค่ maintain กรอบ top-k ไว้ก็พอ ทำให้เหลือ O(n log k) ที่ดีกว่าเมื่อ k น้อย

initialize (ตั้งค่าเริ่มต้น) min-heap ว่าง ๆ ชื่อ heap
iterate (วน) ทีละตัว n ใน array: push n เข้า heap
ถ้าหลัง push แล้ว heap ยาวเกิน k ตัว ให้ heappop ทิ้งตัว minimum (มันเล็กเกินกว่าจะติด top-k)
จบ loop heap เหลือ k ตัวที่มากที่สุดของทั้ง array ตัว minimum ในนั้น (heap[0]) คือคำตอบ

จุดพลาดที่พบบ่อย

สับสนว่าต้องใช้ max-heap แต่จริง ๆ การหา k ตัวมากสุดกลับใช้ min-heap เพราะเราอยากให้ตัวเล็กสุดในกลุ่มถูก evict (เขี่ยออก) ได้ง่าย ๆ อีก edge case คือ k เท่ากับความยาว array ก็จะได้ตัว minimum ของทั้ง array ซึ่งถูกต้อง

ไล่ทีละสเต็ป

จำลอง nums = [3,2,1,5,6,4], k = 2 (maintain ไว้แค่ 2 ตัวมากสุด):

push n	heap หลัง push	ยาวเกิน k?	heap หลังจัดการ
3	[3]	ไม่	[3]
2	[2, 3]	ไม่	[2, 3]
1	[1, 3, 2]	ใช่ pop 1	[2, 3]
5	[2, 3, 5]	ใช่ pop 2	[3, 5]
6	[3, 5, 6]	ใช่ pop 3	[5, 6]
4	[4, 6, 5]	ใช่ pop 4	[5, 6]

จบ loop heap[0] = 5 คือคำตอบ

▶ เฉลยละเอียด (ลองเองก่อนนะ)

เฉลย (Python) — โค้ดนี้รันได้จริงpython

import heapq

def find_kth_largest(nums, k):
    heap = []
    for n in nums:
        heapq.heappush(heap, n)     # ใส่เข้า min-heap
        if len(heap) > k:
            heapq.heappop(heap)     # เกิน k ตัว ทิ้งตัวน้อยสุด
    # เหลือ k ตัวที่มากที่สุด และ heap[0] คือตัวน้อยสุดในนั้น
    return heap[0]

print(find_kth_largest([3, 2, 1, 5, 6, 4], 2))          # 5
print(find_kth_largest([3, 2, 3, 1, 2, 4, 5, 5, 6], 4)) # 4

Output

5
4

ไอเดียคือเราต้องการ maintain k ตัวที่มากที่สุด แล้วในบรรดา k ตัวนั้น ตัว minimum ก็คือคำตอบ (ตัวมากอันดับ k) เราใช้ min-heap ขนาด k เป็นกรอบเก็บ เมื่อ push ตัวใหม่แล้ว heap ยาวเกิน k เราก็ pop ตัว minimum ทิ้งไป (เพราะมันเล็กเกินกว่าจะติด top-k) ตัวที่รอดอยู่จึงเป็น k ตัวใหญ่สุดเสมอ

ถ้าเปลี่ยนไปใช้ max-heap แทน จะกลายเป็นต้อง pop ออก n-k ครั้งเพื่อ access (เข้าถึง) อันดับ k ซึ่งวุ่นกว่า การใช้ min-heap ขนาด k ทำให้ตัวที่เล็กเกินไปหลุดออกเองอัตโนมัติ เหลือแต่ผู้ท้าชิง top-k เท่านั้น

Time O(n log k) iterate ทุกตัว n ครั้ง แต่ละครั้ง push/pop บน heap ขนาด k เป็น O(log k) · Space O(k) heap เก็บอย่างมาก k ตัว

💡 สรุป pattern

โจทย์ top-k ไม่ต้อง sort ทั้ง array แค่ maintain min-heap ขนาด k evict ตัวเล็กสุดออกเรื่อย ๆ ตัวที่รอดคือ top-k และ heap[0] คือตัวอันดับ k พอดี