Trapping Rain Water ใน Python

กำหนดอาร์เรย์ n จำนวนเต็มไม่ติดลบ แทน elevation map ความกว้างแต่ละแท็บ = 1 ต้องคำนวณปริมาณน้ำที่เก็บได้หลังฝนตก

ตัวอย่าง: [0,1,0,2,1,0,1,3,2,1,2,1] → 6 หน่วยน้ำ

ขั้นตอนการแก้ปัญหา (Stack)

• ใช้ stack เก็บ index ของแท็บที่ความสูงลดลง
• วน i: ถ้า stack ว่างหรือ height[stack[-1]] ≥ height[i] → push i, i++
• ถ้า height[stack[-1]] < height[i]: pop x, ถ้า stack ไม่ว่าง คำนวณน้ำ = (min(height[stack[-1]], height[i]) - height[x]) * (i - stack[-1] - 1)

ตัวอย่างโค้ด

class Solution(object):
    def trap(self, height):
        stack = []
        water = 0
        i = 0
        while i < len(height):
            if len(stack) == 0 or height[stack[-1]] >= height[i]:
                stack.append(i)
                i += 1
            else:
                x = stack[-1]
                stack.pop()
                if len(stack) != 0:
                    temp = min(height[stack[-1]], height[i])
                    dist = i - stack[-1] - 1
                    water += dist * (temp - height[x])
        return water

ob = Solution()
print(ob.trap([0,1,0,2,1,0,1,3,2,1,2,1]))

อินพุต

[0,1,0,2,1,0,1,3,2,1,2,1]

เอาต์พุต

6