นากกายมองแท่งความสูงบนโต๊ะมิ้นต์ มีแอ่งน้ำเทลขังระหว่างแท่ง

LeetCode #42 Trapping Rain Water: น้ำขังระหว่างแท่งความสูง

น้องกายEngineering

โจทย์ให้ความสูงแท่งจากซ้ายไปขวา แล้วถามว่าฝนตกเต็มที่แล้วมีน้ำขังกี่หน่วย ตามLeetCode 42 น้ำที่ตำแหน่ง i ขังได้เท่ากับ min(maxซ้าย, maxขวา) - height[i] ถ้าค่านี้บวก

วิธีที่ 1: Two pointers — ค่าเริ่มต้นที่แนะนำ

เดินจากสองด้าน จำ max ซ้าย/ขวา ขยับฝั่งที่ max ต่ำกว่า เวลา O(n) หน่วยความจำ O(1)

วิธีที่ 2: Prefix / suffix max

สร้างอาร์เรย์ max ซ้ายและขวาแล้วรวมน้ำทีละช่อง เวลา O(n) หน่วยความจำ O(n) อ่านแนวคิดง่ายมาก

วิธีที่ 3: Monotonic stack

สแต็กดัชนีสูงลดลง เมื่อเจอแท่งสูงกว่า ให้เอาแอ่งออกทีละชั้น เวลา O(n) ยืดหยุ่นกับโจทย์ช่วงแต่โค้ดยาวกว่า

import unittest

def trap_two_pointers(height):
    if not height:
        return 0
    left, right = 0, len(height) - 1
    left_max = right_max = water = 0
    while left < right:
        if height[left] <= height[right]:
            left_max = max(left_max, height[left])
            water += left_max - height[left]
            left += 1
        else:
            right_max = max(right_max, height[right])
            water += right_max - height[right]
            right -= 1
    return water

def trap_prefix(height):
    n = len(height)
    if n == 0:
        return 0
    left_max = [0] * n
    right_max = [0] * n
    left_max[0] = height[0]
    for i in range(1, n):
        left_max[i] = max(left_max[i - 1], height[i])
    right_max[-1] = height[-1]
    for i in range(n - 2, -1, -1):
        right_max[i] = max(right_max[i + 1], height[i])
    return sum(min(left_max[i], right_max[i]) - height[i] for i in range(n))

def trap_stack(height):
    stack = []
    water = 0
    for i, h in enumerate(height):
        while stack and h > height[stack[-1]]:
            bottom = stack.pop()
            if not stack:
                break
            left = stack[-1]
            width = i - left - 1
            bounded = min(h, height[left]) - height[bottom]
            water += width * bounded
        stack.append(i)
    return water

class TrapTests(unittest.TestCase):
    def test_examples(self):
        cases = [
            ([0, 1, 0, 2, 1, 0, 1, 3, 2, 1, 2, 1], 6),
            ([4, 2, 0, 3, 2, 5], 9),
            ([], 0),
            ([1], 0),
        ]
        for height, expected in cases:
            for method in (trap_two_pointers, trap_prefix, trap_stack):
                with self.subTest(method=method.__name__, height=height):
                    self.assertEqual(method(height), expected)

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

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

เงื่อนไขtwo pointersprefix/suffixstack
หน่วยความจำ O(1)ใช่ไม่ไม่ (O(n))
อธิบายสูตรน้ำขังปานกลางชัดที่สุดชัดเรื่องช่วง
สัมภาษณ์ทั่วไปแนะนำใช้เปิดแนวคิดถ้าผู้สัมภาษณ์ชวน
อินพุตว่าง/แท่งเดียวต้องกันต้องกันต้องกัน

ค่าเริ่มต้นที่ผมใช้: two pointers หลังอธิบายสูตรด้วย prefix ความคิดเพื่อให้คนฟังตามทัน

Recap

  • น้ำขัง = min(maxซ้าย, maxขวา) − ความสูง ณ จุดนั้น
  • two pointers ประหยัดหน่วยความจำ
  • prefix อธิบายง่าย, stack ยืดหยุ่นเรื่องช่วง
  • เทสต์เคสว่าง แท่งเดียว และตัวอย่างทางการ