Count ways to generate N-length array with 0s, 1s, and 2s such that sum of all adjacent pairwise products is K

Given two integers N and K, the task is to find the number of N-length arrays that can be generated by using the values 0, 1, and 2 any number of times, such that the sum of all adjacent pairwise products of the array is K.

Examples:

Input: N = 4, K = 3
Output: 5
Explanation: All possible arrangements are: 

1. arr[] = {2, 1, 1, 0}, Adjacent pairwise product sum = 2 * 1 + 1 * 1 + 1 * 0 = 3.
2. arr[] = {0, 2, 1, 1}, Adjacent pairwise product sum = 0 * 2 + 2 * 1 + 1 * 1 = 3.
3. arr[] = {1, 1, 2, 0}, Adjacent pairwise product sum = 1 * 1 + 1 * 2 + 2 * 0 = 3.
4. arr[] = {0, 1, 1, 2}, Adjacent pairwise product sum is 0 * 1 + 1 * 1 + 1 * 2 = 3.
5. arr[] = {1, 1, 1, 1}, Adjacent pairwise product sum = 1*1 + 1*1 + 1*1 = 3.

Input: N = 10, K = 9
Output: 3445

Naive Approach: The simplest approach is to generate all possible arrangements of the array whose value can be 0, 1, or 2 and count those arrays whose adjacent pairwise product sum is K. Print the count of such arrangements. 

Time Complexity: O(N*3^N )
Auxiliary Space: O(N)

Efficient Approach: To optimize the above approach, the optimal idea is to use Dynamic Programming. The overlapping subproblems can be stored in a dp[][][] table where dp[i][remaining][previous] stores the answer for up to position (N – 1) from position ‘i’ with ‘remaining’ as the remaining value to be added and ‘previous’ as the number placed in the position (i – 1). There can be three cases possible for any position ‘i’:

• Assign ‘0’ to position ‘i’.
• Assign ‘1’ to position ‘i’.
• Assign ‘2’ to position ‘i’.

Follow the steps below to solve the problem:

• Initialize the dp[][][] to store the current position, remaining value to be added, and element at the previous position.
• The transition state is as follows :

dp[i][remaining_sum][previous_element] = dp(assign 0 to pos ‘i’) + dp(assign 1 to ‘i’ ) + dp(assign 2 to ‘i’) 

• Solve the above recurrence relation recursively and store the result for each state in the dp table. For overlapping, subproblems use the stored result in the dp table.
• After the above recursive calls end, print the total number of arrays having adjacent pairwise products of the array is K return by the function.

Below is an implementation of the above approach : 

C++

// C++ program for the above approach

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

// Function to find total number of
// possible arrangements of array
int waysForPairwiseSumToBeK(
    int i, int rem, int previous,
    int N, int dp[][15][3])
{
    // Base Case
    if (i == N) {
        if (rem == 0)
            return 1;
        else
            return 0;
    }

    // If rem exceeds 'k' return 0
    if (rem < 0)
        return 0;

    // Return the already calculated
    // states
    if (dp[i][rem][previous] != -1)
        return dp[i][rem][previous];

    int ways = 0;

    // Place a '0' at current position
    ways += waysForPairwiseSumToBeK(
        i + 1, rem, 0, N, dp);

    // Place a '1' at current position
    // Add it to previous value
    ways += waysForPairwiseSumToBeK(
        i + 1, rem - (previous), 1, N, dp);

    // Place a '2' at current position.
    // Add it to previous value.
    ways += waysForPairwiseSumToBeK(
        i + 1, rem - (2 * previous), 2, N, dp);

    // Store the current state result
    // return the same result
    return dp[i][rem][previous] = ways;
}

// Function to find number of possible
// arrangements of array with 0, 1, and
// 2 having pairwise product sum K
void countOfArrays(int i, int rem,
                   int previous, int N)
{
    // Store the overlapping states
    int dp[15][15][3];

    // Initialize dp table with -1
    memset(dp, -1, sizeof dp);

    // Stores total number of ways
    int totWays
        = waysForPairwiseSumToBeK(
            i, rem, previous, N, dp);

    // Print number of ways
    cout << totWays << ' ';
}

// Driver Code
int main()
{
    // Given N and K
    int N = 4, K = 3;

    // Function Call
    countOfArrays(0, K, 0, N);

    return 0;
}

Java

// Java program for the
// above approach
import java.util.*;
class solution{

// Function to find total number of
// possible arrangements of array
static int waysForPairwiseSumToBeK(int i, int rem,
                                   int previous,
                                   int N, int [][][]dp)
{
  // Base Case
  if (i == N)
  {
    if (rem == 0)
      return 1;
    else
      return 0;
  }

  // If rem exceeds
  // 'k' return 0
  if (rem < 0)
    return 0;

  // Return the already
  // calculated states
  if (dp[i][rem][previous] != -1)
    return dp[i][rem][previous];

  int ways = 0;

  // Place a '0' at current position
  ways += waysForPairwiseSumToBeK(i + 1, rem,
                                  0, N, dp);

  // Place a '1' at current position
  // Add it to previous value
  ways += waysForPairwiseSumToBeK(i + 1, rem -
                                  (previous),
                                  1, N, dp);

  // Place a '2' at current position.
  // Add it to previous value.
  ways += waysForPairwiseSumToBeK(i + 1, rem -
                                  (2 * previous),
                                  2, N, dp);

  // Store the current state result
  // return the same result
  dp[i][rem][previous] = ways;
  return ways;
}

// Function to find number of possible
// arrangements of array with 0, 1, and
// 2 having pairwise product sum K
static void countOfArrays(int i, int rem,
                          int previous, int N)
{
  // Store the overlapping states
  int [][][]dp = new int[15][15][3];

  // Initialize dp table with -1
  for(int p = 0; p < 15; p++)
  {
    for(int q = 0; q < 15; q++)
    {
      for(int r = 0; r < 3; r++)
        dp[p][q][r] = -1;
    }
  }

  // Stores total number of ways
  int totWays = waysForPairwiseSumToBeK(i, rem,
                                        previous,
                                        N, dp);
  // Print number of ways
  System.out.print(totWays);
}

// Driver Code
public static void main(String args[])
{
  // Given N and K
  int N = 4, K = 3;

  // Function Call
  countOfArrays(0, K, 0, N);
}
}

// This code is contributed by SURENDRA_GANGWAR

Python3

# Python3 program for the above approach

# Function to find total number of
# possible arrangements of array
def waysForPairwiseSumToBeK(i, rem, previous, N, dp):

    # Base Case
    if (i == N):
        if (rem == 0):
            return 1
        else:
            return 0

    # If rem exceeds 'k' return 0
    if (rem < 0):
        return 0

    # Return the already calculated
    # states
    if (dp[i][rem][previous] != -1):
        return dp[i][rem][previous]

    ways = 0

    # Place a '0' at current position
    ways += waysForPairwiseSumToBeK(i + 1, rem,
                                    0, N, dp)

    # Place a '1' at current position
    # Add it to previous value
    ways += waysForPairwiseSumToBeK(i + 1,
                                  rem - (previous),
                                  1, N, dp)

    # Place a '2' at current position.
    # Add it to previous value.
    ways += waysForPairwiseSumToBeK(i + 1,
                             rem - (2 * previous),
                             2, N, dp)

    # Store the current state result
    # return the same result
    dp[i][rem][previous] = ways

    return ways

# Function to find number of possible
# arrangements of array with 0, 1, and
# 2 having pairwise product sum K
def countOfArrays(i, rem, previous, N):

    # Store the overlapping states
    dp = [[[-1 for i in range(3)]
               for j in range(15)]
               for k in range(15)]

    # Stores total number of ways
    totWays = waysForPairwiseSumToBeK(i, rem,
                                      previous,
                                      N, dp)

    # Print number of ways
    print(totWays, end = " ")

# Driver Code
if __name__ == '__main__':

    # Given N and K
    N = 4
    K = 3

    # Function Call
    countOfArrays(0, K, 0, N)

# This code is contributed by bgangwar59

C#

// C# program for the
// above approach
using System;

class GFG{

// Function to find total number of
// possible arrangements of array
static int waysForPairwiseSumToBeK(int i, int rem,
                                   int previous,
                                   int N, int [,,]dp)
{

  // Base Case
  if (i == N)
  {
    if (rem == 0)
      return 1;
    else
      return 0;
  }

  // If rem exceeds
  // 'k' return 0
  if (rem < 0)
    return 0;

  // Return the already
  // calculated states
  if (dp[i, rem, previous] != -1)
    return dp[i, rem, previous];

  int ways = 0;

  // Place a '0' at current position
  ways += waysForPairwiseSumToBeK(i + 1, rem,
                                  0, N, dp);

  // Place a '1' at current position
  // Add it to previous value
  ways += waysForPairwiseSumToBeK(i + 1, rem -
                                  (previous),
                                  1, N, dp);

  // Place a '2' at current position.
  // Add it to previous value.
  ways += waysForPairwiseSumToBeK(i + 1, rem -
                                  (2 * previous),
                                   2, N, dp);

  // Store the current state result
  // return the same result
  dp[i, rem, previous] = ways;

  return ways;
}

// Function to find number of possible
// arrangements of array with 0, 1, and
// 2 having pairwise product sum K
static void countOfArrays(int i, int rem,
                          int previous, int N)
{

  // Store the overlapping states
  int [,,]dp = new int[ 15, 15, 3 ];

  // Initialize dp table with -1
  for(int p = 0; p < 15; p++)
  {
    for(int q = 0; q < 15; q++)
    {
      for(int r = 0; r < 3; r++)
        dp[p, q, r] = -1;
    }
  }

  // Stores total number of ways
  int totWays = waysForPairwiseSumToBeK(i, rem,
                                        previous,
                                        N, dp);

  // Print number of ways
  Console.Write(totWays);
}

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

  // Given N and K
  int N = 4, K = 3;

  // Function Call
  countOfArrays(0, K, 0, N);
}
}

// This code is contributed by gauravrajput1

Javascript

<script>
// Javascript program for the
// above approach

// Function to find total number of
// possible arrangements of array
function waysForPairwiseSumToBeK(i,rem,previous,N,dp)
{
    // Base Case
  if (i == N)
  {
    if (rem == 0)
      return 1;
    else
      return 0;
  }

  // If rem exceeds
  // 'k' return 0
  if (rem < 0)
    return 0;

  // Return the already
  // calculated states
  if (dp[i][rem][previous] != -1)
    return dp[i][rem][previous];

  let ways = 0;

  // Place a '0' at current position
  ways += waysForPairwiseSumToBeK(i + 1, rem,
                                  0, N, dp);

  // Place a '1' at current position
  // Add it to previous value
  ways += waysForPairwiseSumToBeK(i + 1, rem -
                                  (previous),
                                  1, N, dp);

  // Place a '2' at current position.
  // Add it to previous value.
  ways += waysForPairwiseSumToBeK(i + 1, rem -
                                  (2 * previous),
                                  2, N, dp);

  // Store the current state result
  // return the same result
  dp[i][rem][previous] = ways;
  return ways;
}

// Function to find number of possible
// arrangements of array with 0, 1, and
// 2 having pairwise product sum K
function countOfArrays(i,rem,previous,N)
{
    // Store the overlapping states
  let dp = new Array(15);
  for(let i = 0; i < 15; i++)
  {
      dp[i] = new Array(15);
    for(let j = 0; j < 15; j++)
    {
        dp[i][j] = new Array(3);
        for(let k = 0; k < 3; k++)
            dp[i][j][k] = -1;
    }
  }

  // Stores total number of ways
  let totWays = waysForPairwiseSumToBeK(i, rem,
                                        previous,
                                        N, dp);
  // Print number of ways
  document.write(totWays);
}

// Driver Code
// Given N and K
let N = 4, K = 3;

// Function Call
countOfArrays(0, K, 0, N);

// This code is contributed by patel2127
</script>

Output:

5

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

Efficient approach : Using DP Tabulation method ( Iterative approach )

The approach to solve this problem is same but DP tabulation(bottom-up) method is better then Dp + memoization(top-down) because memoization method needs extra stack space of recursion calls.

Steps to solve this problem :

• Create a table to store the solution of the subproblems.
• Initialize the table with base cases
• Fill up the table iteratively
• Return the final solution

Implementation :

C++

// C++ program for the above approach

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


// Function to find total number of
// possible arrangements of array
int waysForPairwiseSumToBeK(int N, int K)
{
    int dp[N+1][K+1][3];
    memset(dp, 0, sizeof dp);

    // find solution of current problema through
    // subproblems iteratively
    for(int i = 0; i <= N; i++) {
        for(int rem = 0; rem <= K; rem++) {
            for(int previous = 0; previous <= 2; previous++) {
                // Base case
                if(i == 0) {
                    if(rem == 0) dp[i][rem][previous] = 1;
                    else dp[i][rem][previous] = 0;
                }
                else {
                    int ways = 0;
                    // Place a '0' at current position
                    ways += dp[i-1][rem][0];

                    // Place a '1' at current position
                     // Add it to previous value
                    if(rem - previous >= 0)
                        ways += dp[i-1][rem-previous][1];

                    // Place a '2' at current position.
                    // Add it to previous value.
                    if(rem - (2 * previous) >= 0)
                        ways += dp[i-1][rem-(2*previous)][2];

                      // store computaion of subproblem
                    // in DP and iterate further
                    dp[i][rem][previous] = ways;
                }
            }
        }
    }
    // return answer
    return dp[N][K][0];
}

// Driver code
int main()
{
    int N = 4, K = 3;
    // function call
    int totWays = waysForPairwiseSumToBeK(N, K);
    cout << totWays << ' ';

    return 0;
}

Javascript

// Function to find total number of possible arrangements of array
function waysForPairwiseSumToBeK(N, K) {
    let dp = new Array(N + 1).fill(0).map(() => new Array(K + 1).fill(0).map(() => new Array(3).fill(0)));

    // Find solution of current problem through subproblems iteratively
    for (let i = 0; i <= N; i++) {
        for (let rem = 0; rem <= K; rem++) {
            for (let previous = 0; previous <= 2; previous++) {
                // Base case
                if (i == 0) {
                    if (rem == 0) dp[i][rem][previous] = 1;
                    else dp[i][rem][previous] = 0;
                } else {
                    let ways = 0;
                    // Place a '0' at current position
                    ways += dp[i - 1][rem][0];

                    // Place a '1' at current position
                    // Add it to previous value
                    if (rem - previous >= 0)
                        ways += dp[i - 1][rem - previous][1];

                    // Place a '2' at current position.
                    // Add it to previous value.
                    if (rem - (2 * previous) >= 0)
                        ways += dp[i - 1][rem - (2 * previous)][2];

                    // Store computation of subproblem
                    // in DP and iterate further
                    dp[i][rem][previous] = ways;
                }
            }
        }
    }
    // Return answer
    return dp[N][K][0];
}

let N = 4,
    K = 3;
// Function call
let totWays = waysForPairwiseSumToBeK(N, K);
console.log(totWays);

Python3

# Python program for the above approach

# Function to find total number of
# possible arrangements of array
def waysForPairwiseSumToBeK(N, K):
    dp = [[[0 for _ in range(3)] for _ in range(K+1)] for _ in range(N+1)]

    # find solution of current problema through
    # subproblems iteratively
    for i in range(N+1):
        for rem in range(K+1):
            for previous in range(3):
                # Base case
                if i == 0:
                    if rem == 0:
                        dp[i][rem][previous] = 1
                    else:
                        dp[i][rem][previous] = 0
                else:
                    ways = 0
                    # Place a '0' at current position
                    ways += dp[i-1][rem][0]

                    # Place a '1' at current position
                    # Add it to previous value
                    if rem - previous >= 0:
                        ways += dp[i-1][rem-previous][1]

                    # Place a '2' at current position.
                    # Add it to previous value.
                    if rem - (2 * previous) >= 0:
                        ways += dp[i-1][rem-(2*previous)][2]

                    # store computation of subproblem
                    # in DP and iterate further
                    dp[i][rem][previous] = ways

    # return answer
    return dp[N][K][0]

# Driver code
N = 4
K = 3
# function call
totWays = waysForPairwiseSumToBeK(N, K)
print(totWays)

Output

5 

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