Tìm mảng con từ hai mảng có tổng bằng nhau trong Python

Cho hai mảng P, Q kích thước N chứa các số 1..N. Tìm mảng con của P và mảng con của Q có tổng bằng nhau. Trả về chỉ số.

Ví dụ: P=[2,3,4,5,6], Q=[9,3,2,6,5] → P[0..2]=9, Q[0]=9.

Giải thuật (Tổng tiền tố + hiệu)

• Tính tổng tiền tố cho P và Q.
• Duyệt i trên P, tìm j trên Q sao cho Q[j] ≥ P[i].
• Hiệu diff = Q[j] - P[i]. Nếu diff đã gặp trước → tìm được cặp mảng con.
• Dùng dictionary lưu diff -> (i,j).

Ví dụ

def show_res(x, y, num):
    print(f"Indices of array {num}:", end=" ")
    for i in range(x, y):
        print(i, end=", ")
    print(y)

def get_subarray(P, Q, swap):
    N = len(P)
    index = {0: (-1, -1)}
    diff = 0
    j = 0
    for i in range(N):
        while Q[j] < P[i]:
            j += 1
        diff = Q[j] - P[i]
        if diff in index:
            idx = index[diff]
            if swap:
                show_res(idx[1] + 1, j, 1)
                show_res(idx[0] + 1, i, 2)
            else:
                show_res(idx[0] + 1, i, 1)
                show_res(idx[1] + 1, j, 2)
            return
        index[diff] = (i, j)
    print(-1)

def cumsum(arr):
    for i in range(1, len(arr)):
        arr[i] += arr[i-1]

P = [2, 3, 4, 5, 6]
Q = [9, 3, 2, 6, 5]
cumsum(P)
cumsum(Q)
if Q[-1] > P[-1]:
    get_subarray(P, Q, False)
else:
    get_subarray(Q, P, True)

Đầu vào

[2, 3, 4, 5, 6], [9, 3, 2, 6, 5]

Đầu ra

Indices of array 1 : 0, 1, 2 Indices of array 2 : 0