LeetCode #347 Top K Frequent Elements: กองหรือถังนับความถี่

โจทย์ให้ลิสต์ตัวเลขและจำนวน k แล้วคืนสมาชิกที่โผล่บ่อยสุด k ตัว ตามLeetCode 347 ลำดับในคำตอบมักไม่บังคับ

วิธีที่ 1: Counter + heapq.nlargest — อ่านง่าย

นับความถี่แล้วหยิบ k ค่าที่ความถี่สูงสุดด้วยฮีป เวลาประมาณ O(n log k)

วิธีที่ 2: Bucket ตามความถี่ — เส้นตรงเฉลี่ย

หลังนับความถี่ สร้างถังขนาด n+1 ใส่ค่าตามความถี่ แล้วกวาดจากถังหนาแน่นสุด เวลา O(n) เหมาะเมื่ออยากได้เส้นตรง

วิธีที่ 3: เรียงคู่ (ค่า, ความถี่)

เรียงทั้งตารางความถี่แล้วตัด k ตัวแรก เวลา O(n log n) โค้ดสั้นแต่ไม่จำเป็นเมื่อ n ใหญ่และ k เล็ก

import unittest
from collections import Counter
import heapq

def top_k_heap(nums, k):
    count = Counter(nums)
    return heapq.nlargest(k, count.keys(), key=count.get)

def top_k_bucket(nums, k):
    count = Counter(nums)
    buckets = [[] for _ in range(len(nums) + 1)]
    for num, freq in count.items():
        buckets[freq].append(num)
    ans = []
    for freq in range(len(buckets) - 1, 0, -1):
        for num in buckets[freq]:
            ans.append(num)
            if len(ans) == k:
                return ans
    return ans

def top_k_sort(nums, k):
    count = Counter(nums)
    return [num for num, _ in sorted(count.items(), key=lambda kv: kv[1], reverse=True)[:k]]

class TopKTests(unittest.TestCase):
    def test_examples(self):
        cases = [
            ([1, 1, 1, 2, 2, 3], 2, {1, 2}),
            ([1], 1, {1}),
        ]
        for nums, k, expected in cases:
            for method in (top_k_heap, top_k_bucket, top_k_sort):
                with self.subTest(method=method.__name__, nums=nums, k=k):
                    self.assertEqual(set(method(nums, k)), expected)

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

เลือกตามสถานการณ์

เงื่อนไข	heap	bucket	sort
k << n	ดี	ดี	เกินจำเป็น
อยาก O(n)	ไม่ใช่	ใช่	ไม่ใช่
โค้ดสั้นสัมภาษณ์	ดีมาก	ยาวกว่า	สั้น
จำโครงสร้าง	heapq	list ของ list	sorted

ค่าเริ่มต้นที่ผมใช้ในสัมภาษณ์: heap — ถ้าผู้สัมภาษณ์ขอเส้นตรง ค่อยขยับไป bucket

Recap

• นับความถี่ก่อนเสมอ
• heap เหมาะ k เล็ก, bucket เหมาะอยาก O(n)
• อย่าสมมติลำดับคำตอบถ้าโจทย์ไม่ล็อก
• เทสต์ด้วยเซตเมื่อลำดับไม่สำคัญ