Find Sub-Arrays with Equal Sum in Python

Suppose we have two arrays P and Q whose size are N, they are holding numbers 1 to N. We have to find sub-arrays from the given arrays so that they have equal sum. Finally return the indices of such sub-arrays. If there is no solution, then return -1.

So, if the input is like P = [2, 3, 4, 5, 6], Q = [9, 3, 2, 6, 5], then the output will be Indices in first array : 0, 1, 2 and indices in second array: 0, so P[0..2] = 2 + 3 + 4 = 9 and Q[0] = 9.

To solve this, we will follow these steps −

• Define a function get_subarray() . This will take P, Q, swap

• N := size of P

• index := a new map

• difference := 0, j := 0

• index[0] := a pair like (-1, -1)

• for i in range 0 to N, do

  • while Q[j] < P[i], do

    • j := j + 1

  • difference := Q[j] - P[i]

  • if difference present in index, then

    • if swap is ture, then

      • idx := index[Q[j] - P[i]]

      • display all values from idx[1] + 1 to j for P

      • display all values from idx[0] + 1 to i for Q

    • otherwise,

      • idx := index[Q[j] - P[i]]

      • display all values from idx[0] + 1 to i for P

      • display all values from idx[1] + 1 to j for Q

    • return

  • index[difference] := (i, j)

• display -1

• From the main methood, do the following −

• Update P and Q using their cumulative sums

• N := size of P

• if Q[N - 1] > P[N - 1], then

  • get_subarray(P, Q, False)

• otherwise,

  • get_subarray(Q, P, True)

Example 

Let us see the following implementation to get better understanding −

 Live Demo

def show_res(x, y, num):
   print("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 = {}
   difference, j = 0, 0
   index[0] = (-1, -1)
   for i in range(0, N):
      while Q[j] < P[i]:
         j += 1
      difference = Q[j] - P[i]
      if difference in index:
         if swap:
            idx = index[Q[j] - P[i]]
            show_res(idx[1] + 1, j, 1)
            show_res(idx[0] + 1, i, 2)
         else:
            idx = index[Q[j] - P[i]]
            show_res(idx[0] + 1, i, 1)
            show_res(idx[1] + 1, j, 2)
         return
      index[difference] = (i, j)
   print(-1)
def cumsum(arr):
   n = len(arr)
   for i in range(1, n):
      arr[i] += arr[i - 1]
P = [2, 3, 4, 5, 6]
Q = [9, 3, 2, 6, 5]
cumsum(P)
cumsum(Q)
N = len(P)
if Q[N - 1] > P[N - 1]:
   get_subarray(P, Q, False)
else:
   get_subarray(Q, P, True)

Input

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

Output

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