แบ่งอาร์เรย์ 3 ส่วนผลรวมเท่ากันใน Python

ให้อาร์เรย์ A ของจำนวนเต็ม ตรวจสอบว่าสามารถแบ่งเป็น 3 ส่วนต่อเนื่อง ไม่ว่าง ที่มีผลรวมเท่ากันได้หรือไม่

อย่างเป็นทางการ: หา i+1 < j ที่ทำให้ sum(A[0..i]) == sum(A[i+1..j-1]) == sum(A[j..n-1])

ตัวอย่าง: [0,2,1,-6,6,-7,9,1,2,0,1] → True ส่วน: [0,2,1], [-6,6,-7,9,1], [2,0,1] (ผลรวม = 3)

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

1. คำนวณผลรวมทั้งหมด temp ถ้าหาร 3 ไม่ลงตัว คืน False
2. required_sum = temp / 3
3. สร้าง prefix sum sum_left และ suffix sum sum_right
4. ใช้ two pointers index1 (ซ้าย) และ index2 (ขวา) หาจุดตัดที่ได้ required_sum
5. ถ้าพบ index1 < index2 และไม่ชนกัน คืน True

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

class Solution(object):
    def canThreePartsEqualSum(self, A):
        temp = sum(A)
        if temp % 3 != 0:
            return False
        sum_left = [0] * len(A)
        sum_left[0] = A[0]
        sum_right = [0] * len(A)
        sum_right[-1] = A[-1]
        for i in range(1, len(A)):
            sum_left[i] = A[i] + sum_left[i-1]
        for i in range(len(A)-2, -1, -1):
            sum_right[i] = A[i] + sum_right[i+1]
        required_sum = temp // 3
        index1 = 0
        index2 = len(A) - 1
        while index1 < index2:
            while index1 < index2 and sum_left[index1] != required_sum:
                index1 += 1
            while index2 > index1 and sum_right[index2] != required_sum:
                index2 -= 1
            return index1 < index2 and index1 != index2

ob1 = Solution()
print(ob1.canThreePartsEqualSum([0,2,2,-6,6,-7,9,2,2,0,2]))

อินพุต

[0,2,1,-6,6,-7,9,1,2,0,1]

เอาต์พุต

True

วิธีแก้แบบ Pythonic (Single Pass)

def canThreePartsEqualSum(A):
    total = sum(A)
    if total % 3 != 0:
        return False
    target = total // 3
    cur_sum = 0
    parts = 0
    for i, num in enumerate(A):
        cur_sum += num
        if cur_sum == target:
            parts += 1
            cur_sum = 0
            if parts == 2 and i != len(A) - 1:
                return True
    return False