Minimum cost to convert all elements of a K-size subarray to 0 from given Ternary Array with subarray sum as cost

Given an array arr[] of N integers, where each array element is either 0, 1, or 2, and an integer K, the task is to print the minimum cost needed to convert all the elements of the array to 0s by selecting a subarray of size K and converting any array element of the subarray to 0 in one operation with the cost equal to the sum of elements of the subarray.

Examples:

Input: arr[] = {2, 0, 1, 1, 0, 2}, K = 3
Output: 3
Explanation:
Perform the following steps:

1. Select the subarray over the range [1, 3], and then flip the arr[2](= 1) to 0 at a cost of 2. The array modifies to arr[] = {2, 0, 0, 1,  0, 2}.
2. Select the subarray over the range [1, 3], and then flip the arr[3](= 1) to 0 at a cost of 1. The array modifies to arr[] = {2, 0, 0, 0, 0, 2}.
3. Select the subarray over the range [0, 2], and then flip the arr[0](= 2) to 0 at a cost of 2. The array modifies to arr[] = {0, 0, 0, 0, 0, 2}.
4. Select the subarray over the range [3, 5], and then flip the arr[5](= 2) to 0 at a cost of 2. The array modifies to arr[] = {0, 0, 0, 0, 0, 0}.

Therefore, the total cost needed is (2+1+2+2 = 7). And it is also the minimum cost needed.

Input: arr[] = {1, 0, 2, 1}, K = 4
Output: 7

Approach: The given problem can be solved based on the following observations: 

1. It is optimal to first convert all the elements of the subarray of size K having the minimum cost and the maximum number of 2s to 0.
2. Suppose X is the count of 1s, and Y is the count of 2s in the subarray with minimum cost. Then converting all 2s to 0 first and then all 1s to 0 will give the minimum cost to convert the subarray to 0.
3. The cost of converting all 2s to 0 is:
   • (X+2*Y)+(X+2*Y-2) +(X+2*Y-4)+ … +(X+2*Y-2*(Y-1)) 
     = Y*(X+2*Y) – Y*(0+2*(Y-1))/2 
     = Y*(X+2*Y) – Y*(Y-1)
     = X*Y+2*Y²-Y²+Y = Y² + X*Y + Y 
4. The cost of converting all remaining 1s to 0 is:
   • X+ (X-1)+(X-2)+ … + 1
     = (X*(X+1))/2
5. Therefore, the cost of converting all elements of a subarray with X 1s and Y 2s is equal to:
   •  Y²+X*Y+Y+(X*(X+1))/2
6. Now, the remaining elements can be converted to 0s by taking a subarray with K-1 0s and the element, which will be of cost equal to the element itself.
7. Therefore, if the sum is the total sum of the array, then, the minimum cost will be equal to:
   • sum-(X+2*Y) + Y²+X*Y+Y+(X*(X+1))/2

Follow the steps below to solve the problem:

• Initialize two arrays, say pcount1[] and pcount2[] as {0} where pcount1[i] and pcount2[i] store the count of 1s and 2s in the prefix over the range [0, i].
• Traverse the array, arr[] using the variable i and do the following:
  • Assign pcount1[i] to pcount1[i+1] and pcount2[i] to pcount2[i+1].
  • If arr[i] is 1, then increment pcount1[i+1]. Otherwise, if arr[i] is 2, then the increment count of pcount2[i+1].
• Find the sum of the array and store it in a variable, say sum.
• Initialize two variables, say X and Y, to store the count of 1s and 2s in the subarray with a minimum sum of size K.
• Iterate over the range [K, N] using the variable i and perform the following steps:
  • Initialize two variables, say A and B, to store the count of 1s and 2s in the subarray over the range [i-K, K-1].
  • Assign pcount1[i]-pcount1[i-K] to A, and pcount2[i]-pcount2[i-K] to B.
  • If A+2*B is less than X+2*Y, then update X to A and Y to B.
• Finally, after completing the above steps, print the total cost as, sum-(X+2*Y) + Y²+X*Y+Y+(X*(X+1))/2.

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 the min cost
int minCost(int arr[], int N, int K)
{
    // Stores the prefix count of 1s
    int pcount1[N + 1] = { 0 };
    // Stores the prefix count of 2s
    int pcount2[N + 1] = { 0 };

    // Traverse the array arr[]
    for (int i = 1; i <= N; i++) {
        pcount1[i] = pcount1[i - 1] + (arr[i - 1] == 1);
        pcount2[i] = pcount2[i - 1] + (arr[i - 1] == 2);
    }

    // Stores total sum of the array arr[]
    int sum = pcount1[N] + 2 * pcount2[N];

    // Stores the count of 1s in a subarray
    int X = N;

    // Stores the count of 2s in a subarray
    int Y = N;

    // Iterate over the range [K, N]
    for (int i = K; i <= N; i++) {
        int A = pcount1[i] - pcount1[i - K];
        int B = pcount2[i] - pcount2[i - K];

        // If current subarray sum is less
        // than X+2*Y
        if (A + 2 * B < X + 2 * Y) {
            X = A;
            Y = B;
        }
    }

    // Stores the total cost
    int total = sum - (X + 2 * Y) + Y * Y + X * Y + Y
                + (X * (X + 1)) / 2;
    // Return total
    return total;
}

// Driver Code
int main()
{

    // Input
    int arr[] = { 2, 0, 1, 1, 0, 2 };
    int N = sizeof(arr) / sizeof(arr[0]);
    int K = 3;
    // Function call
    cout << minCost(arr, N, K);

    return 0;
}

Java

// Java Program for the above approach
import java.io.*;
import java.util.Arrays;
class GFG
{

    // Function to find the min cost
    static int minCost(int arr[], int N, int K)
    {

        // Stores the prefix count of 1s
        int pcount1[] = new int[N + 1];

        // Stores the prefix count of 2s
        int pcount2[] = new int[N + 1];
        Arrays.fill(pcount1, 0);
        Arrays.fill(pcount2, 0);

        // Traverse the array arr[]
        for (int i = 1; i <= N; i++) {
            int k = 0;
            int l = 0;
            if (arr[i - 1] == 1)
                k = 1;

            if (arr[i - 1] == 2)
                l = 1;

            pcount1[i] = pcount1[i - 1] + k;
            pcount2[i] = pcount2[i - 1] + l;
        }

        // Stores total sum of the array arr[]
        int sum = pcount1[N] + 2 * pcount2[N];

        // Stores the count of 1s in a subarray
        int X = N;

        // Stores the count of 2s in a subarray
        int Y = N;

        // Iterate over the range [K, N]
        for (int i = K; i <= N; i++) {
            int A = pcount1[i] - pcount1[i - K];
            int B = pcount2[i] - pcount2[i - K];

            // If current subarray sum is less
            // than X+2*Y
            if (A + 2 * B < X + 2 * Y) {
                X = A;
                Y = B;
            }
        }

        // Stores the total cost
        int total = sum - (X + 2 * Y) + Y * Y + X * Y + Y
                    + (X * (X + 1)) / 2;
        // Return total
        return total;
    }

    // Driver Code
    public static void main(String[] args)
    {

        // Input
        int arr[] = { 2, 0, 1, 1, 0, 2 };
        int N = arr.length;
        int K = 3;
        // Function call

        System.out.println(minCost(arr, N, K));

    }
}

      // This code is contributed by Potta Lokesh

Python3

# Py program for the above approach

# Function to find the min cost
def minCost(arr, N, K):

    # Stores the prefix count of 1s
    pcount1 = [0] * (N + 1)

    # Stores the prefix count of 2s
    pcount2 = [0] * (N+1)

    # Traverse the array arr[]
    for i in range(1,N+1):
        pcount1[i] = pcount1[i - 1] + (arr[i - 1] == 1)
        pcount2[i] = pcount2[i - 1] + (arr[i - 1] == 2)

    # Stores total sum of the array arr[]
    sum = pcount1[N] + 2 * pcount2[N]

    # Stores the count of 1s in a subarray
    X = N

    # Stores the count of 2s in a subarray
    Y = N

    # Iterate over the range [K, N]
    for i in range(K, N + 1):
        A = pcount1[i] - pcount1[i - K]
        B = pcount2[i] - pcount2[i - K]

        # If current subarray sum is less
        # than X+2*Y
        if (A + 2 * B < X + 2 * Y):
            X = A
            Y = B

    # Stores the total cost
    total = sum - (X + 2 * Y) + Y * Y + X * Y + Y + (X * (X + 1)) // 2
    # Return total
    return total

# Driver Code
if __name__ == '__main__':

    # Input
    arr= [2, 0, 1, 1, 0, 2]
    N =  len(arr)
    K = 3

    # Function call
    print (minCost(arr, N, K))

# This code is contributed by mohit kumar 29.

C#

// C# program for the above approach
using System;
using System.Collections.Generic;

class GFG{

// Function to find the min cost
static int minCost(int []arr, int N, int K)
{
    // Stores the prefix count of 1s
    int []pcount1 = new int[N + 1];
    Array.Clear(pcount1, 0, N + 1);

    // Stores the prefix count of 2s
    int []pcount2 = new int[N + 1];
    Array.Clear(pcount2, 0, N + 1);

    // Traverse the array arr[]
    for (int i = 1; i <= N; i++) {
        if(arr[i - 1] == 1){
          pcount1[i] = pcount1[i - 1] + 1;
        }
        else
          pcount1[i] = pcount1[i - 1];

         if(arr[i - 1] == 2){
          pcount2[i] = pcount2[i - 1] + 1;
        }
        else
          pcount2[i] = pcount2[i - 1];

    }

    // Stores total sum of the array arr[]
    int sum = pcount1[N] + 2 * pcount2[N];

    // Stores the count of 1s in a subarray
    int X = N;

    // Stores the count of 2s in a subarray
    int Y = N;

    // Iterate over the range [K, N]
    for (int i = K; i <= N; i++) {
        int A = pcount1[i] - pcount1[i - K];
        int B = pcount2[i] - pcount2[i - K];

        // If current subarray sum is less
        // than X+2*Y
        if (A + 2 * B < X + 2 * Y) {
            X = A;
            Y = B;
        }
    }

    // Stores the total cost
    int total = sum - (X + 2 * Y) + Y * Y + X * Y + Y
                + (X * (X + 1)) / 2;
    // Return total
    return total;
}

// Driver Code
public static void Main()
{

    // Input
    int []arr = { 2, 0, 1, 1, 0, 2 };
    int N = arr.Length;
    int K = 3;

    // Function call
    Console.Write(minCost(arr, N, K));
}
}

// This code is contributed by SURENDRA_GANGWAR.

Javascript

// Javascript program for the above approach

// Function to find the min cost
function minCost(arr, N, K)
{

    // Stores the prefix count of 1s
    let pcount1 = new Array(N + 1).fill(0);

    // Stores the prefix count of 2s
    let pcount2 = new Array(N + 1).fill(0);

    // Traverse the array arr[]
    for (let i = 1; i <= N; i++) {
        pcount1[i] = pcount1[i - 1] + (arr[i - 1] == 1);
        pcount2[i] = pcount2[i - 1] + (arr[i - 1] == 2);
    }

    // Stores total sum of the array arr[]
    let sum = pcount1[N] + 2 * pcount2[N];

    // Stores the count of 1s in a subarray
    let X = N;

    // Stores the count of 2s in a subarray
    let Y = N;

    // Iterate over the range [K, N]
    for (let i = K; i <= N; i++) {
        let A = pcount1[i] - pcount1[i - K];
        let B = pcount2[i] - pcount2[i - K];

        // If current subarray sum is less
        // than X+2*Y
        if (A + 2 * B < X + 2 * Y) {
            X = A;
            Y = B;
        }
    }

    // Stores the total cost
    let total = sum - (X + 2 * Y) + Y * Y + X * Y + Y
        + (X * (X + 1)) / 2;

    // Return total
    return total;
}

// Driver Code

// Input
let arr = [2, 0, 1, 1, 0, 2];
let N = arr.length;
let K = 3;

// Function call
document.write(minCost(arr, N, K));

// This code is contributed by gfgking.

Output

7

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