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

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[diff][/diff] 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++ 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 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 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# 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

<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>

6

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