Given an array arr[] of size N, you can perform the following operation on the array any number of times: If arr[i]=0, we can replace it with any integer. Your task is to maximize the number of indices having prefix sum = 0.

Examples:

Input: N = 7, arr[] = {1, -1, 0, 2, 4, 2, 0, 1}
Output: 3
Explanation: We can have at max 3 indices with prefix sum = 0 by changing arr[6] to -8. So, arr[] will become {1, -1, 0, 2, 4, 2, -8, 1}. Now, the prefix sum array would be {1, 0, 0, 2, 6, 8, 0, 1}

Input: N = 3, arr[] = {1000, 1000, 0}
Output: 1
Explanation: We can have at max 1 index with prefix sum = 0 by changing arr[2] to -2000. So, arr[] will become {1000, 1000, -2000}. Now, the prefix sum array would be {1000, 1000, 0}.

Approach: To solve the problem, follow the below idea:

Let’s consider that the indices where arr[i] = 0 are: z1, z2, z3 …. zx, so we can divide the array arr[] into sub-segments like: [0…z1-1], [z1+1…z2-1], [z2+1…z3-1], and so on. After dividing it into sub-segments,

• For the first segment [0…z1-1], no change can be made to all the prefix sums from [0…z1-1] as we cannot change any element before index z1. So, for segment arr[0…z1-1], we can find the number of indices with prefix sum = 0 by iterating from 0 to z1 and maintaining the frequency of 0 prefix sums for each index.
• Now for all the remaining segments, we can turn the most frequent prefix sum in that segment, say P to 0 by changing the 0 which is just before the segment to -P. So, we can find the number of indices with prefix sum = 0 by maintaining the frequency of prefix sum in that segment and adding it to the answer.

Step-by-step algorithm:

• Maintain a map, say freqMap to store the frequency of prefix sums for every segment.
• Also, maintain a variable maxFreq to store the frequency of the most frequent prefix sum.
• Iterate over the array arr[], and check,
  • If arr[i] = 0, if this is the first occurrence of 0, simply add freqMap[0] to the answer otherwise add maxFreq to the answer. Also, clear the freqMap and reset maxFreq to 0.
  • If arr[i] != 0, add the current element to the prefix sum. Increment the frequency of new prefix sum by 1 and update maxFreq if the new prefix sum becomes the most frequent prefix sum.
• After iterating over the entire array, return the final answer.

Below is the implementation of the algorithm:

#include <bits/stdc++.h>
using namespace std;

// function call to maximize the number of indices with
// prefix sum = 0
int solve(int* arr, int N)
{
    // flag to check for the first occurrence of zero
    bool flag = false;

    // frequency map to track frequency of prefix sums in
    // every segment
    map<long long, long long> freqMap;

    // to track maximum frequency in any segment
    long long maxFreq = 0;

    // to track the prefix sum till any index
    long long prefSum = 0;

    // variable to store the answer
    long long ans = 0;

    for (int i = 0; i < N; i++) {
        if (arr[i] == 0) {

            // If this is first occurrence of zero then we
            // cannot change any prefix sum before inndex i,
            // so we simply add the frequency when we had
            // the prefix sum = 0
            if (flag == false)
                ans += freqMap[0];
            // Else we change the last occurrence of zero in
            // subarray [0...i] to (most frequent element *
            // -1) to get the maximum prefix sum as 0
            else
                ans += maxFreq;

            // mark the occurrence of zero
            flag = true;

            // reset maxFreq to zero and clear the frequency
            // map
            maxFreq = 0;
            freqMap.clear();
        }

        // maintain the prefix Sum
        prefSum += arr[i];

        // Increase the frequency of current prefix sum by 1
        freqMap[prefSum] += 1;

        // Update maxFreq with the frequency of most
        // frequent element
        maxFreq = max(maxFreq, freqMap[prefSum]);
    }

    if (flag == false) {
        ans += freqMap[0];
    }
    else {
        ans += maxFreq;
    }

    return ans;
}

int main()
{
    // Input
    int N = 7;
    int arr[] = { 1, -1, 0, 2, 4, 2, 0, 1 };

    cout << solve(arr, N) << endl;
    return 0;
}

import java.util.HashMap;
import java.util.Map;

class Main {

    // Function to maximize the number of indices with prefix sum = 0
    static long solve(int[] arr, int N) {
        // flag to check for the first occurrence of zero
        boolean flag = false;

        // frequency map to track frequency of prefix sums in every segment
        Map<Long, Long> freqMap = new HashMap<>();

        // to track maximum frequency in any segment
        long maxFreq = 0;

        // to track the prefix sum till any index
        long prefSum = 0;

        // variable to store the answer
        long ans = 0;

        for (int i = 0; i < N; i++) {
            if (arr[i] == 0) {
                // If this is the first occurrence of zero then we cannot change any prefix sum
                // before index i, so we simply add the frequency when we had the prefix sum = 0
                if (!flag)
                    ans += freqMap.getOrDefault(0L, 0L);
                // Else we change the last occurrence of zero in subarray [0...i] to (most frequent element * -1)
                // to get the maximum prefix sum as 0
                else
                    ans += maxFreq;

                // mark the occurrence of zero
                flag = true;

                // reset maxFreq to zero and clear the frequency map
                maxFreq = 0;
                freqMap.clear();
            }

            // maintain the prefix sum
            prefSum += arr[i];

            // Increase the frequency of the current prefix sum by 1
            freqMap.put(prefSum, freqMap.getOrDefault(prefSum, 0L) + 1);

            // Update maxFreq with the frequency of the most frequent element
            maxFreq = Math.max(maxFreq, freqMap.get(prefSum));
        }

        if (!flag) {
            ans += freqMap.getOrDefault(0L, 0L);
        } else {
            ans += maxFreq;
        }

        return ans;
    }

    public static void main(String[] args) {
        // Input
        int N = 7;
        int[] arr = {1, -1, 0, 2, 4, 2, 0, 1};

        System.out.println(solve(arr, N));
    }
}

// This code is contributed by shivamgupta310570

def solve(arr):
    # Initialize variables
    flag = False  # Flag to check for the first occurrence of zero
    freq_map = {}  # Frequency map to track prefix sum frequencies
    max_freq = 0   # Maximum frequency in any segment
    pref_sum = 0   # Prefix sum till any index
    ans = 0        # Variable to store the answer

    # Loop through the array
    for num in arr:
        # Check if the current number is zero
        if num == 0:
            # If this is the first occurrence of zero
            if not flag:
                ans += freq_map.get(0, 0)  # Add the frequency of prefix sum = 0
            else:
                ans += max_freq  # Add the frequency of the most frequent element

            flag = True  # Mark the occurrence of zero
            max_freq = 0  # Reset max_freq
            freq_map = {}  # Clear the frequency map

        # Maintain the prefix sum
        pref_sum += num

        # Increase the frequency of the current prefix sum by 1
        freq_map[pref_sum] = freq_map.get(pref_sum, 0) + 1

        # Update max_freq with the frequency of the most frequent element
        max_freq = max(max_freq, freq_map[pref_sum])

    # If zero was not encountered in the array
    if not flag:
        ans += freq_map.get(0, 0)  # Add the frequency of prefix sum = 0
    else:
        ans += max_freq  # Add the frequency of the most frequent element

    return ans

def main():
    # Input array
    arr = [1, -1, 0, 2, 4, 2, 0, 1]

    # Call solve function and print the result
    print(solve(arr))

if __name__ == "__main__":
    main()

using System;
using System.Collections.Generic;

public class Program
{
    // Function to maximize the number of indices with prefix sum = 0
    static long Solve(int[] arr, int N)
    {
        // Flag to check for the first occurrence of zero
        bool flag = false;

        // Frequency map to track frequency of prefix sums in every segment
        Dictionary<long, long> freqMap = new Dictionary<long, long>();

        // To track maximum frequency in any segment
        long maxFreq = 0;

        // To track the prefix sum till any index
        long prefSum = 0;

        // Variable to store the answer
        long ans = 0;

        for (int i = 0; i < N; i++)
        {
            if (arr[i] == 0)
            {
                // If this is the first occurrence of zero, then we
                // cannot change any prefix sum before index i,
                // so we simply add the frequency when we had
                // the prefix sum = 0
                if (!flag)
                    ans += freqMap.ContainsKey(0) ? freqMap[0] : 0;
                // Else, we change the last occurrence of zero in
                // subarray [0...i] to (most frequent element *
                // -1) to get the maximum prefix sum as 0
                else
                    ans += maxFreq;

                // Mark the occurrence of zero
                flag = true;

                // Reset maxFreq to zero and clear the frequency map
                maxFreq = 0;
                freqMap.Clear();
            }

            // Maintain the prefix sum
            prefSum += arr[i];

            // Increase the frequency of the current prefix sum by 1
            if (freqMap.ContainsKey(prefSum))
                freqMap[prefSum]++;
            else
                freqMap[prefSum] = 1;

            // Update maxFreq with the frequency of the most frequent element
            maxFreq = Math.Max(maxFreq, freqMap[prefSum]);
        }

        if (!flag)
            ans += freqMap.ContainsKey(0) ? freqMap[0] : 0;
        else
            ans += maxFreq;

        return ans;
    }

    // Driver Code
    public static void Main()
    {
        // Input
        int N = 7;
        int[] arr = { 1, -1, 0, 2, 4, 2, 0, 1 };

        Console.WriteLine(Solve(arr, N));
    }
}

// This code is contributed by shivamgupta310570

// Function to maximize the number of indices with prefix sum = 0
function solve(arr) {
    // Flag to check for the first occurrence of zero
    let flag = false;

    // Map to track frequency of prefix sums in every segment
    const freqMap = new Map();

    // Variable to track maximum frequency in any segment
    let maxFreq = 0;

    // Variable to track the prefix sum till any index
    let prefSum = 0;

    // Variable to store the answer
    let ans = 0;

    for (let i = 0; i < arr.length; i++) {
        if (arr[i] === 0) {
            // If this is the first occurrence of zero, add the frequency
            // when the prefix sum was 0
            if (!flag)
                ans += freqMap.get(0) || 0;
            // Else, change the last occurrence of zero in the subarray [0...i]
            // to (most frequent element * -1) to get the maximum prefix sum as 0
            else
                ans += maxFreq;

            // Mark the occurrence of zero
            flag = true;

            // Reset maxFreq to zero and clear the frequency map
            maxFreq = 0;
            freqMap.clear();
        }

        // Maintain the prefix sum
        prefSum += arr[i];

        // Increase the frequency of current prefix sum by 1
        freqMap.set(prefSum, (freqMap.get(prefSum) || 0) + 1);

        // Update maxFreq with the frequency of the most frequent element
        maxFreq = Math.max(maxFreq, freqMap.get(prefSum) || 0);
    }

    // If zero was not found, add the frequency when the prefix sum was 0
    if (!flag) {
        ans += freqMap.get(0) || 0;
    } else {
        ans += maxFreq;
    }

    return ans;
}

// Input
const arr = [1, -1, 0, 2, 4, 2, 0, 1];

console.log(solve(arr)); // Output: 5

3

Time Complexity: O(N*logN), where N is the size of the input array arr[].
Auxiliary Space: O(N)