Check if it possible to 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

Arnab Chakraborty

Published on 19-Jan-2021 04:34:24

• Related Questions & Answers
• Partition to K Equal Sum Subsets in C++
• Program to check whether we can partition a list with k-partitions of equal sum in C++
• Partition Array Into Three Parts With Equal Sum in Python
• JavaScript Total subarrays with Sum K
• Partition Equal Subset Sum in C++
• Finding n subarrays with equal sum in JavaScript
• Check if it is possible to survive on Island in Python
• Check if it is possible to serve customer queue with different notes in Python
• Check if it is possible to sort the array after rotating it in Python
• Check if array sum can be made K by three operations on it in Python
• Check if it is possible to create a polygon with a given angle in Python
• Check if it is possible to create a polygon with given n sidess in Python
• Check if it is possible to transform one string to another in Python
• Check if it is possible to convert one string into another with given constraints in Python
• Find all subarrays with sum equal to number? JavaScript (Sliding Window Algorithm)