Count ways to choose elements from given Arrays

Given arrays A[] and B[] of size N, the task is to count ways to form an array of size N such that the new array formed has i^th element either A[i] or B[i] and adjacent elements are different.

Examples:

Input: A[] = {1, 4, 3}, B[] = {2, 2, 4}
Output: 4
Explanation: Possible arrays are {1, 4, 3}, {1, 2, 3}, {1, 2, 4} and {2, 4, 3}.
Below is the Recursion Tree. Green is a valid move whereas red is an invalid move.
 

Input: A[] = {1, 2, 3, 4}, B[] = {5, 6, 7, 8}
Output:  16

Naive approach: The basic way to solve the problem is as follows:

The basic way to solve this problem is to generate all possible combinations by using a recursive approach.

Time Complexity: O(2^N)
Auxiliary Space: O(1)

Efficient Approach/Memoization:  The above approach can be optimized based on the following idea:

Dynamic programming can be used to solve this problem:

• dp[i][j] represents number of possible arrays of size i with last element chosen from array j (0 for A[] and 1 for B[]).

It can be observed that the recursive function is called exponential times. That means that some states are called repeatedly. 
So the idea is to store the value of each state. This can be done by storing the value of a state and whenever the function is called, returning the stored value without computing again.

Follow the steps below to solve the problem:

• Declare a dp array, dp[100001][2] initialized with -1 using memset() function.
• Create a 2d array Arr[][], whose first column will contain array A[] and the second column will have array B[].
• Create a recursive function that takes four parameters i and j such that i represents position and j represents the last element chosen from either A[] or B[], Arr[] stores array A[] and B[], and size of array N.
• Calling recursive function recur(0, 0, Arr, N) with i = 0 indicates we will fill the first position of the required array.
• Check if i == N, return 1.
• Check if dp[i][j] != -1, return dp[i][j].
• Initialize variable ans = 0 to count the number of valid solutions.
• Call recursive function for both choosing an element from A[] (in this case Arr[][0]) or B[] (in this case Arr[][1]) and add the answer returned to cnt .
• Return ans calculated for the current state.
• Save the answer in dp[i][j] before returning.

Below is the implementation of the above approach:

// C++ code to implement the approach
#include <bits/stdc++.h>
using namespace std;

// mod to avoid integer overflow
const int mod = 1e9 + 7;

// dp table
int dp[100001][2];

// recursive Function to count ways to
// form array of size N such that adjcent
// elements chosen are different
int recur(int i, int j, vector<vector<int> >& Arr, int N)
{

    // If already computed state just
    // return dp[i][j]
    if (dp[i][j] != -1)
        return dp[i][j];

    // If array of size N is possible
    if (i == N)
        return 1;

    // Answer initialized with value zero
    int ans = 0;

    // if choosing element from A[]
    if (i == 0 or Arr[i - 1][j] != Arr[i][0])
        ans = (ans + recur(i + 1, 0, Arr, N)) % mod;

    // if choosing element from B[]
    if (i == 0 or Arr[i - 1][j] != Arr[i][1])
        ans = (ans + recur(i + 1, 1, Arr, N)) % mod;

    // Return the final answer
    return dp[i][j] = ans;
}

// Function to count ways to form
// array of size N such that adjcent
// elements chosen are different
void findWays(int A[], int B[], int N)
{

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

    // making one array for A[] & B[]
    vector<vector<int> > Arr(N, vector<int>(2));

    // filling Arr
    for (int i = 0; i < N; i++) {
        Arr[i][0] = A[i];
        Arr[i][1] = B[i];
    }

    // calling recursive function
    cout << recur(0, 0, Arr, N) << endl;
}

// Driver Code
int main()
{

    // Input 1
    int A[] = { 1, 4, 3 }, B[] = { 2, 2, 4 };
    int N = sizeof(A) / sizeof(A[0]);

    // Function Call
    findWays(A, B, N);

    // Input 2
    int A1[] = { 1, 2, 3, 4 }, B1[] = { 5, 6, 7, 8 };
    int N1 = sizeof(A1) / sizeof(A1[0]);

    // Function Call
    findWays(A1, B1, N1);

    return 0;
}

// Java code to implement the approach

import java.io.*;
import java.util.*;

class GFG {

    // mod to avoid integer overflow
    static final int mod = (int)1e9 + 7;

    // dp table
    static int[][] dp = new int[100001][2];

    // recursive Function to count ways to form array of
    // size N such that adjacent elements chosen are
    // different
    static int recur(int i, int j, int[][] Arr, int N)
    {

        // If already computed state just return dp[i][j]
        if (dp[i][j] != -1)
            return dp[i][j];

        // If array of size N is possible
        if (i == N)
            return 1;

        // Answer initialized with value zero
        int ans = 0;

        // if choosing element from A[]
        if (i == 0 || Arr[i - 1][j] != Arr[i][0])
            ans = (ans + recur(i + 1, 0, Arr, N)) % mod;

        // if choosing element from B[]
        if (i == 0 || Arr[i - 1][j] != Arr[i][1])
            ans = (ans + recur(i + 1, 1, Arr, N)) % mod;

        // Return the final answer
        return dp[i][j] = ans;
    }

    // Function to count ways to form array of size N such
    // that adjacent elements chosen are different
    static void findWays(int[] A, int[] B, int N)
    {

        // initializing dp table with -1
        for (int i = 0; i <= N; i++) {
            Arrays.fill(dp[i], -1);
        }

        // making one array for A[] & B[]
        int[][] Arr = new int[N][2];

        // filling Arr
        for (int i = 0; i < N; i++) {
            Arr[i][0] = A[i];
            Arr[i][1] = B[i];
        }

        // calling recursive function
        System.out.println(recur(0, 0, Arr, N));
    }

    public static void main(String[] args)
    {
        // Input 1
        int[] A = { 1, 4, 3 };
        int[] B = { 2, 2, 4 };
        int N = A.length;

        // Function Call
        findWays(A, B, N);

        // Input 2
        int[] A1 = { 1, 2, 3, 4 };
        int[] B1 = { 5, 6, 7, 8 };
        int N1 = A1.length;

        // Function Call
        findWays(A1, B1, N1);
    }
}

// This code is contributed by lokesh.

# python code to implement the approach
MOD = int(1e9) + 7

# dp table
dp = [[-1 for j in range(2)] for i in range(100001)]

# recursive Function to count ways to
# form array of size N such that adjacent
# elements chosen are different


def recur(i, j, arr, N):
    global dp

    # If already computed state just
    # return dp[i][j]
    if dp[i][j] != -1:
        return dp[i][j]

    # If array of size N is possible
    if i == N:
        return 1

    # Answer initialized with value zero
    ans = 0

    # if choosing element from A[]
    if i == 0 or arr[i - 1][j] != arr[i][0]:
        ans = (ans + recur(i + 1, 0, arr, N)) % MOD

    # if choosing element from B[]
    if i == 0 or arr[i - 1][j] != arr[i][1]:
        ans = (ans + recur(i + 1, 1, arr, N)) % MOD

    # Return the final answer
    dp[i][j] = ans
    return ans

# Function to count ways to form
# array of size N such that adjacent
# elements chosen are different


def findWays(A, B, N):
    global dp

    # making one array for A[] & B[]
    arr = [[0 for j in range(2)] for i in range(N)]

    # filling arr
    for i in range(N):
        arr[i][0] = A[i]
        arr[i][1] = B[i]

    # initializing dp table with -1
    for i in range(100001):
        dp[i][0] = dp[i][1] = -1

    # calling recursive function
    print(recur(0, 0, arr, N))


# Driver Code
if __name__ == '__main__':

    # Input 1
    A = [1, 4, 3]
    B = [2, 2, 4]
    N = len(A)

    # Function Call
    findWays(A, B, N)

    # Input 2
    A1 = [1, 2, 3, 4]
    B1 = [5, 6, 7, 8]
    N1 = len(A1)

    # Function Call
    findWays(A1, B1, N1)
# This code is contributed by lokesh.

// C# code for the approach

using System;

public class GFG {
    // mod to avoid integer overflow
    const int mod = 1000000007;

    // dp table
    static int[, ] dp = new int[100001, 2];

    // recursive Function to count ways to
    // form array of size N such that adjcent
    // elements chosen are different
    static int recur(int i, int j, int[][] Arr, int N)
    {

        // If already computed state just
        // return dp[i][j]
        if (dp[i, j] != -1)
            return dp[i, j];

        // If array of size N is possible
        if (i == N)
            return 1;

        // Answer initialized with value zero
        int ans = 0;

        // if choosing element from A[]
        if (i == 0 || Arr[i - 1][j] != Arr[i][0])
            ans = (ans + recur(i + 1, 0, Arr, N)) % mod;

        // if choosing element from B[]
        if (i == 0 || Arr[i - 1][j] != Arr[i][1])
            ans = (ans + recur(i + 1, 1, Arr, N)) % mod;

        // Return the final answer
        return dp[i, j] = ans;
    }

    // Function to count ways to form
    // array of size N such that adjcent
    // elements chosen are different
    static void findWays(int[] A, int[] B, int N)
    {

        // initializing dp table with - 1
        for (int i = 0; i < 100001; i++)
            for (int j = 0; j < 2; j++)
                dp[i, j] = -1;

        // making one array for A[] & B[]
        int[][] Arr = new int[N][];

        // filling Arr
        for (int i = 0; i < N; i++) {
            Arr[i] = new int[2];
            Arr[i][0] = A[i];
            Arr[i][1] = B[i];
        }

        // calling recursive function
        Console.WriteLine(recur(0, 0, Arr, N));
    }

    // Driver Code
    public static void Main()
    {

        // Input 1
        int[] A = { 1, 4, 3 }, B = { 2, 2, 4 };
        int N = A.Length;

        // Function Call
        findWays(A, B, N);

        // Input 2
        int[] A1 = { 1, 2, 3, 4 }, B1 = { 5, 6, 7, 8 };
        int N1 = A1.Length;

        // Function Call
        findWays(A1, B1, N1);
    }
}

// mod to avoid integer overflow
const mod = 1e9 + 7;

// dp table
const dp = new Array(100001).fill(null).map(() => new Array(2).fill(-1));

// recursive Function to count ways to
// form array of size N such that adjacent
// elements chosen are different
function recur(i, j, Arr, N) {

  // If already computed state just
  // return dp[i][j]
  if (dp[i][j] !== -1) {
    return dp[i][j];
  }

  // If array of size N is possible
  if (i === N) {
    return 1;
  }

  // Answer initialized with value zero
  let ans = 0;

  // if choosing element from A[]
  if (i === 0 || Arr[i - 1][j] !== Arr[i][0]) {
    ans = (ans + recur(i + 1, 0, Arr, N)) % mod;
  }

  // if choosing element from B[]
  if (i === 0 || Arr[i - 1][j] !== Arr[i][1]) {
    ans = (ans + recur(i + 1, 1, Arr, N)) % mod;
  }

  // Return the final answer
  return dp[i][j] = ans;
}

// Function to count ways to form
// array of size N such that adjacent
// elements chosen are different
function findWays(A, B, N) {

  for(let i=0; i<100001; i++){
      for(let j=0; j<2; j++)
        dp[i][j]=-1;
  }
  // making one array for A[] & B[]
  const Arr = new Array(N).fill(null).map((_, i) => [A[i], B[i]]);

  // calling recursive function
  console.log(recur(0, 0, Arr, N));
}

// Driver Code
(function main() {

  // Input 1
  const A = [1, 4, 3], B = [2, 2, 4];
  const N = A.length;

  // Function Call
  findWays(A, B, N);

  // Input 2
  const A1 = [1, 2, 3, 4], B1 = [5, 6, 7, 8];
  const N1 = A1.length;

  // Function Call
  findWays(A1, B1, N1);

})();

4
16

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

Related Articles:

• Introduction to Dynamic Programming – Data Structures and Algorithms Tutorials