Count ways to split array into two equal sum subarrays by replacing each array element to 0 once

0 0 9 minutes read

Given an array arr[] consisting of N integers, the task is to count the number of ways to split the array into two subarrays of equal sum after changing a single array element to 0.

Examples:

Input: arr[] = {1, 2, -1, 3}
Output: 4
Explanation:
Replacing arr[0] by 0, arr[] is modified to {0, 2, -1, 3}. Only 1 possible split is {0, 2} and {-1, 3}.
Replacing arr[1] by 0, arr[] is modified to {1, 0, -1, 3}. No way to split the array.
Replacing arr[2] by 0, arr[] is modified to {1, 2, 0, 3}. The 2 possible splits are {1, 2, 0} and {3}, {1, 2} and {0, 3}.
Replacing arr[3] by 0, arr[] is modified to {1, 2, -1, 0}. Only 1 possible split is {1} and {2, -1, 0}.
Therefore, the total number of ways to split = 1 + 0 + 2 + 1 = 4.

Input: arr[] = {1, 2, 1, 1, 3, 1}
Output: 6
Explanation:
Replacing arr[0] by 0, arr[] is modified to {0, 2, 1, 1, 3, 1}. Only 1 possible split exists.
Replacing arr[1] by 0, arr[] is modified to {1, 0, 1, 1, 3, 1}. No way to split the array.
Replacing arr[2] by 0, arr[] is modified to {1, 2, 0, 1, 3, 1}. Only 1 possible split exists.
Replacing arr[3] by 0, arr[] is modified to {1, 2, 1, 0, 3, 1}. Only 2 possible splits exist.
Replacing arr[4] by 0, arr[] is modified to {1, 2, 1, 1, 0, 1}. Only 1 possible split exists.
Replacing arr[5] by 0, arr[] is modified to {1, 2, 1, 1, 3, 0}. Only 1 possible split exists.
Total number of ways to split is = 1 + 0 + 1 + 2 + 1 + 1 = 6.

Naive Approach: The simplest approach to solve the problem is to traverse the array, convert each array element arr[i] to 0 and count the number of ways to split the modified array into two subarrays with equal sum.

Time Complexity: O(N²)
Auxiliary Space: O(1)

Efficient Approach: To optimize the above approach, the idea is based on the following observations:

Considering two arrays arr1[] and arr2[] with sum of the array elements equal to sum_arr1 and sum_arr2 respectively.
Let dif be the difference between sum_arr1 and sum_arr2, i.e., sum_arr1 – sum_arr2 = dif.
Now, sum_arr1 can be made equal to sum_arr2 by performing any one of the two operations:

Remove an element from arr1[] whose value is equal to dif.

Remove an element from arr2[] whose value is equal to -dif.

Therefore, the total number of ways to obtain sum_arr1 = sum_arr2 is equal to the count of dif in arr1 + count of (-dif) in arr2.

For every index in the range [0, N – 1], the total number of ways can be obtained by considering the current index as the splitting point, by making any element equal to 0 using the process discussed above. Follow the steps below to solve the problem:

Initialize a variable count with 0 to store the desired result and prefix_sum the with 0 to store the prefix sum and suffixSum with 0 to store the suffix sum.
Initialize hashmaps prefixCount and suffixCount to store the count of elements in prefix and suffix arrays.
Traverse the arr[] and update the frequency of each element in suffixCount.
Traverse the arr[] over the range [0, N – 1] using variable i.
- Add arr[i] to the prefixCount hashmap and remove it from suffixCount.
- Add arr[i] to prefixSum and set suffixSum to the difference of the total sum of the array and prefixSum.
- Store the difference between the sum of a subarray in variable dif = prefix_sum – suffixSum.
- Store the number of ways to split at i^th index in number_of_subarray_at_i_split and is equal to the sum of prefixCount and suffixCount.
- Update the count by adding number_of_subarray_at_i_split to it.
After the above steps, print the value of count as the result.

Below is the implementation of the above approach:

C++

// C++ program for the above approach
#include <bits/stdc++.h>
using namespace std;
 
// Function to find number of ways to
// split array into 2 subarrays having
// equal sum by changing element to 0 once
int countSubArrayRemove(int arr[], int N)
{
    // Stores the count of elements
    // in prefix and suffix of
    // array elements
    unordered_map<int, int>
        prefix_element_count,
        suffix_element_count;
 
    // Stores the sum of array
    int total_sum_of_elements = 0;
 
    // Traverse the array
    for (int i = N - 1; i >= 0; i--) {
 
        total_sum_of_elements += arr[i];
 
        // Increase the frequency of
        // current element in suffix
        suffix_element_count[arr[i]]++;
    }
 
    // Stores prefix sum upto index i
    int prefix_sum = 0;
 
    // Stores sum of suffix of index i
    int suffix_sum = 0;
 
    // Stores the desired result
    int count_subarray_equal_sum = 0;
 
    // Traverse the array
    for (int i = 0; i < N; i++) {
 
        // Modify prefix sum
        prefix_sum += arr[i];
 
        // Add arr[i] to prefix map
        prefix_element_count[arr[i]]++;
 
        // Calculate suffix sum by
        // subtracting prefix sum
        // from total sum of elements
        suffix_sum = total_sum_of_elements
                     - prefix_sum;
 
        // Remove arr[i] from suffix map
        suffix_element_count[arr[i]]--;
 
        // Store the difference
        // between the subarrays
        int difference = prefix_sum
                         - suffix_sum;
 
        // Count number of ways to split
        // the array at index i such that
        // subarray sums are equal
        int number_of_subarray_at_i_split
            = prefix_element_count[difference]
              + suffix_element_count[-difference];
 
        // Update the final result
        count_subarray_equal_sum
            += number_of_subarray_at_i_split;
    }
 
    // Return the result
    return count_subarray_equal_sum;
}
 
// Driver Code
int main()
{
    int arr[] = { 1, 2, 1, 1, 3, 1 };
    int N = sizeof(arr) / sizeof(arr[0]);
 
    // Function Call
    cout << countSubArrayRemove(arr, N);
 
    return 0;
}

Java

// Java program for the above approach
import java.util.*; 
 
class GFG{
     
// Function to find number of ways to
// split array into 2 subarrays having
// equal sum by changing element to 0 once
static int countSubArrayRemove(int []arr, int N)
{
     
    // Stores the count of elements
    // in prefix and suffix of
    // array elements
    HashMap<Integer,
            Integer> prefix_element_count = new HashMap<Integer,
                                                        Integer>();
    HashMap<Integer,
            Integer> suffix_element_count = new HashMap<Integer, 
                                                        Integer>();
 
    // Stores the sum of array
    int total_sum_of_elements = 0;
 
    // Traverse the array
    for(int i = N - 1; i >= 0; i--)
    {
        total_sum_of_elements += arr[i];
 
        // Increase the frequency of
        // current element in suffix
        if (!suffix_element_count.containsKey(arr[i]))
            suffix_element_count.put(arr[i], 1);
        else
            suffix_element_count.put(arr[i],
              suffix_element_count.get(arr[i]) + 1);
    }
 
    // Stores prefix sum upto index i
    int prefix_sum = 0;
 
    // Stores sum of suffix of index i
    int suffix_sum = 0;
 
    // Stores the desired result
    int count_subarray_equal_sum = 0;
 
    // Traverse the array
    for(int i = 0; i < N; i++)
    {
         
        // Modify prefix sum
        prefix_sum += arr[i];
 
        // Add arr[i] to prefix map
       if (!prefix_element_count.containsKey(arr[i]))
           prefix_element_count.put(arr[i], 1);
       else
           prefix_element_count.put(arr[i],
           prefix_element_count.get(arr[i]) + 1);
 
        // Calculate suffix sum by
        // subtracting prefix sum
        // from total sum of elements
        suffix_sum = total_sum_of_elements - 
                     prefix_sum;
 
        // Remove arr[i] from suffix map
        if (!suffix_element_count.containsKey(arr[i]))
            suffix_element_count.put(arr[i], 0);
        else
            suffix_element_count.put(arr[i],
            suffix_element_count.get(arr[i]) - 1);
 
        // Store the difference
        // between the subarrays
        int difference = prefix_sum - 
                         suffix_sum;
 
        // Count number of ways to split
        // the array at index i such that
        // subarray sums are equal
        int number_of_subarray_at_i_split = 0;
        if (prefix_element_count.containsKey(difference))
            number_of_subarray_at_i_split =
            prefix_element_count.get(difference);
        if (suffix_element_count.containsKey(-difference))
            number_of_subarray_at_i_split += 
            suffix_element_count.get(-difference);
 
        // Update the final result
        count_subarray_equal_sum += 
        number_of_subarray_at_i_split;
    }
 
    // Return the result
    return count_subarray_equal_sum;
}   
 
// Driver Code
public static void main(String args[]) 
{
    int []arr = { 1, 2, 1, 1, 3, 1 };
    int N = arr.length;
     
    // Function Call
    System.out.println(countSubArrayRemove(arr, N));
}
}
 
// This code is contributed by Stream_Cipher

Python3

# Python3 program for the above approach
 
# Function to find number of ways to
# split array into 2 subarrays having
# equal sum by changing element to 0 once
def countSubArrayRemove(arr, N):
     
    # Stores the count of elements
    # in prefix and suffix of
    # array elements
    prefix_element_count = {}
    suffix_element_count = {}
 
    # Stores the sum of array
    total_sum_of_elements = 0
 
    # Traverse the array
    i = N - 1
    while (i >= 0):
        total_sum_of_elements += arr[i]
 
        # Increase the frequency of
        # current element in suffix
        suffix_element_count[arr[i]] = suffix_element_count.get(
            arr[i], 0) + 1
             
        i -= 1
 
    # Stores prefix sum upto index i
    prefix_sum = 0
 
    # Stores sum of suffix of index i
    suffix_sum = 0
 
    # Stores the desired result
    count_subarray_equal_sum = 0
 
    # Traverse the array
    for i in range(N):
         
        # Modify prefix sum
        prefix_sum += arr[i]
 
        # Add arr[i] to prefix map
        prefix_element_count[arr[i]] = prefix_element_count.get(
            arr[i], 0) + 1
 
        # Calculate suffix sum by
        # subtracting prefix sum
        # from total sum of elements
        suffix_sum = total_sum_of_elements - prefix_sum
 
        # Remove arr[i] from suffix map
        suffix_element_count[arr[i]] = suffix_element_count.get(
            arr[i], 0) - 1
 
        # Store the difference
        # between the subarrays
        difference = prefix_sum - suffix_sum
 
        # Count number of ways to split
        # the array at index i such that
        # subarray sums are equal
        number_of_subarray_at_i_split = (prefix_element_count.get(
                                             difference, 0) +
                                         suffix_element_count.get(
                                            -difference, 0))
 
        # Update the final result
        count_subarray_equal_sum += number_of_subarray_at_i_split
 
    # Return the result
    return count_subarray_equal_sum
 
# Driver Code
if __name__ == '__main__':
     
    arr = [ 1, 2, 1, 1, 3, 1 ]
    N =  len(arr)
 
    # Function Call
    print(countSubArrayRemove(arr, N))
     
# This code is contributed by SURENDRA_GANGWAR

C#

// C# program for the above approach
using System;
using System.Collections.Generic;
 
class GFG{
     
// Function to find number of ways to
// split array into 2 subarrays having
// equal sum by changing element to 0 once
static int countSubArrayRemove(int []arr, int N)
{
     
    // Stores the count of elements
    // in prefix and suffix of
    // array elements
    Dictionary<int, 
               int> prefix_element_count = new Dictionary<int,
                                                          int> ();
    Dictionary<int,
               int >suffix_element_count = new Dictionary <int,
                                                           int>();
 
    // Stores the sum of array
    int total_sum_of_elements = 0;
 
    // Traverse the array
    for(int i = N - 1; i >= 0; i--)
    {
        total_sum_of_elements += arr[i];
 
        // Increase the frequency of
        // current element in suffix
        if (!suffix_element_count.ContainsKey(arr[i]))
            suffix_element_count[arr[i]] = 1;
        else
            suffix_element_count[arr[i]]++;
    }
 
    // Stores prefix sum upto index i
    int prefix_sum = 0;
 
    // Stores sum of suffix of index i
    int suffix_sum = 0;
 
    // Stores the desired result
    int count_subarray_equal_sum = 0;
 
    // Traverse the array
    for(int i = 0; i < N; i++)
    {
         
        // Modify prefix sum
        prefix_sum += arr[i];
 
        // Add arr[i] to prefix map
       if (!prefix_element_count.ContainsKey(arr[i]))
           prefix_element_count[arr[i]] = 1;
       else
           prefix_element_count[arr[i]]++;
 
        // Calculate suffix sum by
        // subtracting prefix sum
        // from total sum of elements
        suffix_sum = total_sum_of_elements - 
                     prefix_sum;
 
        // Remove arr[i] from suffix map
        if (!suffix_element_count.ContainsKey(arr[i]))
            suffix_element_count[arr[i]] = 0;
        else
            suffix_element_count[arr[i]]-= 1;
 
        // Store the difference
        // between the subarrays
        int difference = prefix_sum - 
                         suffix_sum;
 
        // Count number of ways to split
        // the array at index i such that
        // subarray sums are equal
        int number_of_subarray_at_i_split = 0;
        if (prefix_element_count.ContainsKey(difference))
            number_of_subarray_at_i_split
            = prefix_element_count[difference];
        if (suffix_element_count.ContainsKey(-difference))
            number_of_subarray_at_i_split 
            += suffix_element_count[-difference];
 
        // Update the final result
        count_subarray_equal_sum
            += number_of_subarray_at_i_split;
    }
 
    // Return the result
    return count_subarray_equal_sum;
}
 
// Driver Code
public static void Main(string []args)
{
    int []arr = { 1, 2, 1, 1, 3, 1 };
    int N = arr.Length;
 
    // Function Call
    Console.Write(countSubArrayRemove(arr, N));
}
}
 
// This code is contributed by chitranayal

Javascript

<script>
 
// Javascript program for the above approach
 
// Function to find number of ways to
// split array into 2 subarrays having
// equal sum by changing element to 0 once
function countSubArrayRemove(arr, N)
{
     
    // Stores the count of elements
    // in prefix and suffix of
    // array elements
    let prefix_element_count = new Map();
    let suffix_element_count = new Map();
  
    // Stores the sum of array
    let total_sum_of_elements = 0;
  
    // Traverse the array
    for(let i = N - 1; i >= 0; i--)
    {
        total_sum_of_elements += arr[i];
  
        // Increase the frequency of
        // current element in suffix
        if (!suffix_element_count.has(arr[i]))
            suffix_element_count.set(arr[i], 1);
        else
            suffix_element_count.set(arr[i],
            suffix_element_count.get(arr[i]) + 1);
    }
  
    // Stores prefix sum upto index i
    let prefix_sum = 0;
  
    // Stores sum of suffix of index i
    let suffix_sum = 0;
  
    // Stores the desired result
    let count_subarray_equal_sum = 0;
  
    // Traverse the array
    for(let i = 0; i < N; i++)
    {
         
        // Modify prefix sum
        prefix_sum += arr[i];
  
        // Add arr[i] to prefix map
       if (!prefix_element_count.has(arr[i]))
           prefix_element_count.set(arr[i], 1);
       else
           prefix_element_count.set(arr[i],
           prefix_element_count.get(arr[i]) + 1);
  
        // Calculate suffix sum by
        // subtracting prefix sum
        // from total sum of elements
        suffix_sum = total_sum_of_elements -
                     prefix_sum;
  
        // Remove arr[i] from suffix map
        if (!suffix_element_count.has(arr[i]))
            suffix_element_count.set(arr[i], 0);
        else
            suffix_element_count.set(arr[i],
            suffix_element_count.get(arr[i]) - 1);
  
        // Store the difference
        // between the subarrays
        let difference = prefix_sum -
                         suffix_sum;
  
        // Count number of ways to split
        // the array at index i such that
        // subarray sums are equal
        let number_of_subarray_at_i_split = 0;
        if (prefix_element_count.has(difference))
            number_of_subarray_at_i_split =
            prefix_element_count.get(difference);
        if (suffix_element_count.has(-difference))
            number_of_subarray_at_i_split +=
            suffix_element_count.get(-difference);
  
        // Update the final result
        count_subarray_equal_sum +=
        number_of_subarray_at_i_split;
    }
  
    // Return the result
    return count_subarray_equal_sum;
}
 
// Driver Code
let arr = [ 1, 2, 1, 1, 3, 1 ];
let  N = arr.length;
 
// Function Call
document.write(countSubArrayRemove(arr, N));
 
// This code is contributed by avanitrachhadiya2155
 
</script>

Output:

Time Complexity: O(N)
Auxiliary Space: O(N)

Feeling lost in the world of random DSA topics, wasting time without progress? It’s time for a change! Join our DSA course, where we’ll guide you on an exciting journey to master DSA efficiently and on schedule.
Ready to dive in? Explore our Free Demo Content and join our DSA course, trusted by over 100,000 zambiatek!