แท่ง prefix sum บนโต๊ะมิ้นต์ ลูกศรชี้ช่วงที่ผลต่างเป็น k สีอำพัน

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 กับอินพุตติดลบโดยไม่พิสูจน์