LeetCode #560 Subarray Sum Equals K: prefix sum + ฮัชแมปนับครั้ง
น้องกายEngineering
Subarray Sum Equals K นับจำนวนช่วงต่อเนื่องที่ผลรวมเป็น k — มีติดลบได้ครับ
วิธี 1: ลองทุกคู่ i≤j
O(n²) ใช้สอนนิยาม
วิธี 2: prefix sum + dict
เก็บจำนวนครั้งที่เคยเห็น prefix — ถ้า prefix - k เคยมี แปลว่ามี subarray จบตรงนี้เท่ากับ k
วิธี 3: sliding window
ใช้ได้เฉพาะเมื่อทุกตัวไม่ติดลบ — โจทย์นี้ทั่วไปใช้ไม่ได้
from collections import defaultdict
import unittest
def subarray_sum(nums: list[int], k: int) -> int:
counts: dict[int, int] = defaultdict(int)
counts[0] = 1
prefix = total = 0
for value in nums:
prefix += value
total += counts[prefix - k]
counts[prefix] += 1
return total
class SubarraySumTests(unittest.TestCase):
def test_examples(self):
self.assertEqual(subarray_sum([1, 1, 1], 2), 2)
self.assertEqual(subarray_sum([1, 2, 3], 3), 2)
self.assertEqual(subarray_sum([1, -1, 0], 0), 3)
if __name__ == "__main__":
unittest.main()
จุดพลาดบ่อย: ลืม counts[0] = 1 สำหรับช่วงที่เริ่มจากหัวลิสต์
Recap
- วาด prefix แล้วหา
prefix - k - เทสต์เคสมีศูนย์และติดลบ
- อย่าใช้ sliding window กับอินพุตติดลบโดยไม่พิสูจน์