Chương trình đếm số cách chia mảng thành ba mảng con trong Python

Giả sử ta có mảng nums, cần tìm số cách chia tốt (good ways) để tách mảng này. Kết quả modulo 10^9 + 7. Một cách chia là tốt nếu mảng được tách thành ba mảng con không rỗng, liên tiếp từ trái sang phải, và tổng mảng trái ≤ tổng mảng giữa ≤ tổng mảng phải.

Ví dụ, nums = [2,3,3,3,7,1], đầu ra là 3 vì có 3 cách chia:

• [2], [3], [3,3,7,1]
• [2], [3,3], [3,7,1]
• [2,3], [3,3], [7,1]

Giải pháp

Để giải quyết bài toán này, ta sử dụng tiền tố tổng (prefix sum) và kỹ thuật two pointers:

• n := kích thước nums, m := 10^9+7
• ss := mảng kích thước n+1 điền 0
• Với mỗi index i, value val trong nums (bắt đầu từ 1): ss[i] := ss[i-1] + val
• r := 0, rr := 0, ans := 0
• Với l từ 1 đến n-2:
  • r := max(r, l+1)
  • Khi r < n-1 và ss[r] - ss[l] < ss[l]: r := r + 1
  • rr := max(rr, r)
  • Khi rr < n-1 và ss[n] - ss[rr+1] >= ss[rr+1] - ss[l]: rr := rr + 1
  • Nếu ss[l] > ss[r] - ss[l]: break
  • Nếu ss[r] - ss[l] > ss[n] - ss[r]: continue
  • ans := (ans + rr - r + 1) % m
• Trả về ans

Ví dụ minh họa

Xem triển khai sau để hiểu rõ hơn:

def solve(nums):
    n, m = len(nums), 10**9+7
    ss = [0] * (1+n)
    for i, val in enumerate(nums, 1):
        ss[i] = ss[i-1] + val
    r = rr = ans = 0
    for l in range(1, n-1):
        r = max(r, l+1)
        while r < n-1 and ss[r]-ss[l] < ss[l]:
            r += 1
        rr = max(rr, r)
        while rr < n-1 and ss[n]-ss[rr+1] >= ss[rr+1]-ss[l]:
            rr += 1
        if ss[l] > ss[r]-ss[l]:
            break
        if ss[r]-ss[l] > ss[n]-ss[r]:
            continue
        ans = (ans+rr-r+1) % m
    return ans

nums = [2,3,3,3,7,1]
print(solve(nums))

Đầu vào

[1,7,3,6,5]

Đầu ra

3