Find partitions that maximises sum of count of 0’s in left part and count of 1’s in right part

Given a binary array nums of length N, The task is to find all possible partitions in given array such that sum of count of 0’s in left part and count of 1’s in right part is maximized.  

Example:

Input: nums = {0, 0, 1, 0}
Output: 2, 4
Explanation: Division at index
index – 0: numsleft is []. numsright is [0, 0, 1, 0]. The sum is 0 + 1 = 1.
index – 1: numsleft is [0]. numsright is [0, 1, 0]. The sum is 1 + 1 = 2.
index – 2: numsleft is [0, 0]. numsright is [1, 0]. The sum is 2 + 1 = 3.
index – 3: numsleft is [0, 0, 1]. numsright is [0]. The sum is 2 + 0 = 2.
index – 4: numsleft is [0, 0, 1, 0]. numsright is []. The sum is 3 + 0 = 3.
Therefore partition at index 2 and 4 gives the highest possible sum 3.

Input: nums= {0, 0, 0, 1, 0, 1, 1}
Output: 3, 5
Explanation: 
If we divide the array at index 3 then numsLeft is 3 and numsRight is 3. The sum is 3 + 3 = 6
If we divide the array at index 5 then numsLeft is 4 and numsRight is 2. The sum is 4 + 2 = 6
Any other division will result in score less than 6.

Approach: The idea is to use prefix sum such that sum[i+1] will be A[0] + … + A[i]. 

• Find prefix sum of the array and store it in array sum.
• For each index i, there will be  i – sum[i] zeros in the left part and sum[N] – sum[i] ones in the right part .
• Calculate score by summing the left part and right part.
• Compare the score with previous max score
  • if current score is greater than the previous one then update the max score and put i in the result array .
  • if current score is equal to  the previous max one then in previous result array push this index i.
• Return the resultant array after complete traversal.

Below is the implementation of above approach:

C++

// C++ program to find partitions
// that maximizes sum of count of 0's
// in left part and count
// of 1's in right part
#include <bits/stdc++.h>
using namespace std;

// Function to find all possible
// max sum partitions
vector<int> findPartitions(vector<int>& A)
{
    int N = A.size(), mx = 0;

    // Prefix sum array
    vector<int> sum(N + 1, 0), ans;
    for (int i = 0; i < N; ++i)
        sum[i + 1] = sum[i] + A[i];

    // Find the partitions
    // with maximum prefix sum
    for (int i = 0; i <= N; ++i) {

        // Calculate score for
        // the current index
        int score = i - sum[i]
                    + sum[N] - sum[i];

        // Compare current score
        // with previous max
        if (score > mx) {
            ans = { i };
            mx = score;
        }

        // If current score
        // is greater than
        // previous then push
        // in the ans array.
        else if (score == mx)
            ans.push_back(i);
    }

    // Returning the resultant
    // array of indices
    return ans;
}

// Driver code
int main()
{
    vector<int> res,
        arr{ 0, 0, 0, 1, 0, 1, 1 };
    res = findPartitions(arr);
    for (auto it : res)
        cout << it << " ";
    return 0;
}

Java

// Java program to find partitions
// that maximizes sum of count of 0's
// in left part and count
// of 1's in right part
import java.util.*;
class GFG{

  // Function to find all possible
  // max sum partitions
  static Vector<Integer>findPartitions(int []A)
  {
    int N = A.length, mx = 0;

    // Prefix sum array
    Vector<Integer> ans = new Vector<Integer>();
    int []sum = new int[N + 1];
    for (int i = 0; i < N; ++i)
      sum[i + 1] = sum[i] + A[i];

    // Find the partitions
    // with maximum prefix sum
    for (int i = 0; i <= N; ++i) {

      // Calculate score for
      // the current index
      int score = i - sum[i]
        + sum[N] - sum[i];

      // Compare current score
      // with previous max
      if (score > mx) {
        ans = new Vector<Integer>();
        ans.add(i);
        mx = score;
      }

      // If current score
      // is greater than
      // previous then push
      // in the ans array.
      else if (score == mx)
        ans.add(i);
    }

    // Returning the resultant
    // array of indices
    return ans;
  }

  // Driver code
  public static void main(String[] args)
  {
    int[] arr = { 0, 0, 0, 1, 0, 1, 1 };
    Vector<Integer> res = findPartitions(arr);
    for (int it : res)
      System.out.print(it+ " ");
  }
}

// This code is contributed by shikhasingrajput

Python3

# Python program to find partitions
# that maximizes sum of count of 0's
# in left part and count
# of 1's in right part

# Function to find all possible
# max sum partitions
def findPartitions(A):

    N = len(A)
    mx = 0

    # Prefix sum array
    sum = []
    for i in range(0, N + 1):
        sum.append(0)
    ans = []
    for i in range(0, N):
        sum[i + 1] = sum[i] + A[i];

    # Find the partitions
    # with maximum prefix sum
    for i in range(0, N + 1):

        # Calculate score for
        # the current index
        score = i - sum[i] + sum[N] - sum[i]

        # Compare current score
        # with previous max
        if (score > mx):
            ans.clear()
            ans.append(i)
            mx = score

        # If current score
        # is greater than
        # previous then push
        # in the ans array.
        elif (score == mx):
            ans.append(i)

    # Returning the resultant
    # array of indices
    return ans

# Driver code
res = []
arr = [ 0, 0, 0, 1, 0, 1, 1 ]
res = findPartitions(arr)
print(res)

# This code is contributed by Samim Hossain Mondal

C#

// C# program to find partitions
// that maximizes sum of count of 0's
// in left part and count
// of 1's in right part
using System;
using System.Collections;

class GFG
{
// Function to find all possible
// max sum partitions
static ArrayList findPartitions(ArrayList A)
{
    int N = A.Count, mx = 0;

    // Prefix sum array
    int []sum = new int[N + 1];
    for(int i = 0; i < N + 1; i++) {
      sum[i] = 0;
    }

    ArrayList ans = new ArrayList();
    for (int i = 0; i < N; ++i)
        sum[i + 1] = sum[i] + (int)A[i];

    // Find the partitions
    // with maximum prefix sum
    for (int i = 0; i <= N; ++i) {

        // Calculate score for
        // the current index
        int score = i - sum[i]
                    + sum[N] - sum[i];

        // Compare current score
        // with previous max
        if (score > mx) {
            ans.Clear();
            ans.Add(i);
            mx = score;
        }

        // If current score
        // is greater than
        // previous then push
        // in the ans array.
        else if (score == mx)
            ans.Add(i);
    }

    // Returning the resultant
    // array of indices
    return ans;
}

// Driver code
public static void Main()
{
    ArrayList arr = new ArrayList();
    arr.Add(0);
    arr.Add(0);
    arr.Add(0);
    arr.Add(1);
    arr.Add(0);
    arr.Add(1);
    arr.Add(1);

    ArrayList res = findPartitions(arr);
    for(int i = 0; i < res.Count; i++)
        Console.Write(res[i] + " ");
}
}

// This code is contributed by Samim Hossain Mondal.

Javascript

<script>
        // JavaScript code for the above approach

        // Function to find all possible
        // max sum partitions
        function findPartitions(A) {
            let N = A.length, mx = 0;

            // Prefix sum array
            let sum = new Array(N + 1).fill(0), ans = [];
            for (let i = 0; i < N; ++i)
                sum[i + 1] = sum[i] + A[i];

            // Find the partitions
            // with maximum prefix sum
            for (let i = 0; i <= N; ++i) {

                // Calculate score for
                // the current index
                let score = i - sum[i]
                    + sum[N] - sum[i];

                // Compare current score
                // with previous max
                if (score > mx) {
                    ans = [i];
                    mx = score;
                }

                // If current score
                // is greater than
                // previous then push
                // in the ans array.
                else if (score == mx)
                    ans.push(i);
            }

            // Returning the resultant
            // array of indices
            return ans;
        }

        // Driver code

        let res,
            arr = [0, 0, 0, 1, 0, 1, 1];
        res = findPartitions(arr);
        for (let it of res)
            document.write(it + " ");

       // This code is contributed by Potta Lokesh
    </script>

Output

3 5 

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