Partition in K Subarrays with Equal Sum in Python

Suppose we have an array of numbers called nums, and also have another value K. We have to check whether we can partition the array nums into K contiguous subarrays such that the elements sum of each subarrays is equal.

So, if the input is like nums = [2, 5, 3, 4, 7] k = 3, then the output will be True as we can make three partitions like [(2, 5), (3, 4), (7)] all have equal sum 7.

To solve this, we will follow these steps −

• n := size of nums
• cumul_sum := cumulative sum of all elements in nums
• total_sum := cumul_sum[n - 1]
• if total_sum not divisible by k, then
  • return False
• count := 0, pos := -1
• for i in range 0 to n - 1, do
  • if pos is same as -1, then
    • sub := 0
  • otherwise,
    • sub := cumul_sum[pos]
  • if cumul_sum[i] - sub is same as (total_sum / K), then
    • pos := i
    • count := count + 1
  • otherwise when cumul_sum[i] - cumul_sum[pos] > (total_sum / K), then
    • come out from loop
• return true when count is same as K otherwise false

Example

Let us see the following implementation to get better understanding −

 Live Demo

def solve(nums, k):
   n = len(nums)
   cumul_sum = [0 for i in range(n)]
   cumul_sum[0] = nums[0]
   for i in range(1, n):
      cumul_sum[i] = cumul_sum[i - 1] + nums[i]
   total_sum = cumul_sum[n - 1]
   if total_sum % k != 0:
      return False
   count = 0
   pos = -1
   for i in range(n):
      if pos == -1:
         sub = 0
      else:
         sub = cumul_sum[pos]
      if cumul_sum[i] - sub == total_sum / k:
         pos = i
         count += 1
      elif cumul_sum[i] - cumul_sum[pos] > total_sum / k:
         break
   return count == k
nums = [2, 5, 3, 4, 7]
k = 3
print(solve(nums, k))

Input

[2, 5, 3, 4, 7], 3

Output

True