Count of ways to select prime or non-prime number based on Array index

Given an array A[] of size N. An array has to be created using the given array considering the following conditions.

• If the index is prime, you must choose a non-prime number that is less than or equal to A[i].
• If the index is non-prime, you must choose a prime number that it less than or equal to A[i].

The task is to count the total number of ways such numbers can be selected.

Note: The indexing of the given array should be considered 1-based indexing.

Examples:

Input: N = 5  A = {2, 3, 4, 8, 5}
Output:  16
Explanation:  You can choose 1 number for index 1 i.e., 2 
(As index 1 is not prime and prime number count less than 
or equal to 2 is one i.e. 2), 1 number for index 2,  
2 numbers for index 3, 4 numbers for index 4 
and 2 numbers for index 5. 
Hence total number of ways = 1x1x2x4x2 = 16.

Input: N = 2  A = {5, 6}
Output:  9
Explanation:  You can choose 3 number for index 1,  
3 numbers for index 2. Hence total number of ways = 3×3 = 9 .

Approach: The idea to solve the problem is as follows:

• Counting and store the value of all non-prime and prime in an array till the maximum element of the array. 
• Then iterate the given array and then if index i is non-prime we multiply the prime count till A[i] and perform the similar operation for prime index.

Follow the below illustration for a better understanding

Illustration: 

Consider an example N = 5 and A[] = {2, 3, 4, 8, 5}

As index 1 is Non Prime So Prime number count less than or equal to 2 is 1 (i.e 2) 
As index 2 is Prime So Non Prime number count less than or equal to 3 is 1 (i.e 1)
As index 3 is Prime So Non Prime number count less than or equal to 4 is 2 (i.e 1, 4)
As index 4 is Non Prime So Prime number count less than or equal to 8 is 4 (i.e 2, 3, 5, 7)
As index 5 is Prime So Non Prime number count less than or equal to 5 is 2 (i.e 1, 4)

Total Number of ways = 1 x 1 x 2 x 4 x 2 = 16

Hence Total number of ways to select number from array is 16.      

Follow the steps mentioned below to implement the idea

• Find the maximum number from the given array.
• Iterate from 1 to the maximum value and find the count of primes and non-primes till every value and store them in a vector of pairs.
• Iterate over the array:
  • Check if the current index is prime or nonprime. if the current index is prime then select the non-prime value count from the vector of the pair.
  • Multiply the answer with the non-prime count and store these values in the answer again.
  • If the current index is non-prime then select prime value count from the vector of pair and multiply with the answer and store these values in answer again.
• Return the answer.

Below is the implementation of the above approach:

C++

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

// Function to check whether the number
// is prime or not
bool isPrime(int a)
{
    if (a == 1)
        return false;
    if (a == 2)
        return true;
    for (int i = 2; i <= sqrt(a); i++) {
        if (a % i == 0)
            return false;
    }
    return true;
}

// Function to count prime and non prime number
// and push count in vector
void count_prime_and_NonPrime(int n,
                              vector<pair<int, int> >& v)
{
    int p = 0, np = 0;
    v.push_back(make_pair(p, np));
    for (int i = 1; i <= n; i++) {
        if (isPrime(i))
            p++;
        else
            np++;
        v.push_back(make_pair(p, np));
    }
}

// Function to find number of ways
int NoOfWays(int n, int a[])
{
    vector<pair<int, int> > v;
    int mx = 0;
    for (int i = 0; i < n; i++) {
        mx = max(mx, a[i]);
    }
    count_prime_and_NonPrime(mx, v);
    int ans = 1;
    for (int j = 0; j < n; j++) {
        int prime = v[a[j]].first;
        int nonPrime = v[a[j]].second;
        if (isPrime(j + 1)) {
            ans *= nonPrime;
        }
        else {
            ans *= prime;
        }
    }
    return ans;
}

// Driver code
int main()
{
    int N = 5;
    int A[] = { 2, 3, 4, 8, 5 };

    // Function call
    cout << NoOfWays(N, A) << endl;
    return 0;
}

Java

// Java program for the above approach

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

class GFG {

    static class pair {
        int first, second;
        public pair(int first, int second)
        {
            this.first = first;
            this.second = second;
        }
    }

    // Function to check whether the number is prime or not
    static boolean isPrime(int a)
    {
        if (a == 1) {
            return false;
        }
        if (a == 2) {
            return true;
        }
        for (int i = 2; i <= Math.sqrt(a); i++) {
            if (a % i == 0) {
                return false;
            }
        }
        return true;
    }

    // Function to count prime and non prime number and push
    // count in arraylist.
    static List<pair> count_prime_and_NonPrime(int n,
                                               List<pair> v)
    {
        int p = 0, np = 0;
        v.add(new pair(p, np));
        for (int i = 1; i <= n; i++) {
            if (isPrime(i)) {
                p++;
            }
            else {
                np++;
            }
            v.add(new pair(p, np));
        }
        return v;
    }

    // Function to find number of ways
    static int NoOfWays(int n, int[] a)
    {
        List<pair> v = new ArrayList<pair>();
        int mx = 0;
        for (int i = 0; i < n; i++) {
            mx = Math.max(mx, a[i]);
        }
        v = count_prime_and_NonPrime(mx, v);
        int ans = 1;
        for (int j = 0; j < n; j++) {
            int prime = v.get(a[j]).first;
            int nonPrime = v.get(a[j]).second;
            if (isPrime(j + 1)) {
                ans *= nonPrime;
            }
            else {
                ans *= prime;
            }
        }
        return ans;
    }

    public static void main(String[] args)
    {
        int N = 5;
        int[] A = { 2, 3, 4, 8, 5 };

        // Function call
        System.out.println(NoOfWays(N, A));
    }
}

// This code is contributed by lokeshmvs21.

Python3

# Function to check whether the number is prime or not
def isPrime(a):
    if (a == 1):
        return False
    if (a == 2):
        return True
    for i in range(2, int(a**0.5)+1):
        if (a % i == 0):
            return False
    return True

# Function to count prime and non prime number and push
# count in arraylist.
def count_prime_and_NonPrime(n, v):
    p = 0
    np = 0
    v = []
    v.append((p, np))
    for i in range(1, n+1):
        if (isPrime(i)):
            p += 1
        else:
            np += 1
        v.append((p, np))
    return v

# Function to find number of ways
def NoOfWays(n, a):
    v = []
    mx = 0
    for i in range(n):
        mx = max(mx, a[i])
    v = count_prime_and_NonPrime(mx, v)
    ans = 1
    for j in range(n):
        prime = v[a[j]][0]
        nonPrime = v[a[j]][1]
        if (isPrime(j + 1)):
            ans *= nonPrime
        else:
            ans *= prime
    return ans


if __name__ == '__main__':
    A = [2, 3, 4, 8, 5]
    N = len(A)
    # Function call
    print(NoOfWays(N, A))

    # This code is contributed by vikkycirus.

C#

// C# program for the above approach
using System;
using System.Collections.Generic;
public class GFG{

  class pair {
    public int first, second;
    public pair(int first, int second)
    {
      this.first = first;
      this.second = second;
    }
  }

  // Function to check whether the number is prime or not
  static bool isPrime(int a)
  {
    if (a == 1) {
      return false;
    }
    if (a == 2) {
      return true;
    }
    for (int i = 2; i <= Math.Sqrt(a); i++) {
      if (a % i == 0) {
        return false;
      }
    }
    return true;
  }

  // Function to count prime and non prime number and push
  // count in arraylist.
  static List<pair> count_prime_and_NonPrime(int n,
                                             List<pair> v)
  {
    int p = 0, np = 0;
    v.Add(new pair(p, np));
    for (int i = 1; i <= n; i++) {
      if (isPrime(i)) {
        p++;
      }
      else {
        np++;
      }
      v.Add(new pair(p, np));
    }
    return v;
  }

  // Function to find number of ways
  static int NoOfWays(int n, int[] a)
  {
    List<pair> v = new List<pair>();
    int mx = 0;
    for (int i = 0; i < n; i++) {
      mx = Math.Max(mx, a[i]);
    }
    v = count_prime_and_NonPrime(mx, v);
    int ans = 1;
    for (int j = 0; j < n; j++) {
      int prime = v[a[j]].first;
      int nonPrime = v[a[j]].second;
      if (isPrime(j + 1)) {
        ans *= nonPrime;
      }
      else {
        ans *= prime;
      }
    }
    return ans;
  }

  static public void Main ()
  {
    int N = 5;
    int[] A = { 2, 3, 4, 8, 5 };

    // Function call
    Console.Write(NoOfWays(N, A));
  }
}

// This code is contributed by hrithikgarg03188.

Javascript

// Javascript program for the above approach

class pair {
    constructor(first, second) {
        this.first = first;
        this.second = second;
    }
}

// Function to check whether the number is prime or not
function isPrime(a) {
    if (a == 1) {
        return false;
    }
    if (a == 2) {
        return true;
    }
    for (let i = 2; i <= Math.floor(Math.sqrt(a)); i++) {
        if (a % i == 0) {
            return false;
        }
    }
    return true;
}

// Function to count prime and non prime number and push
// count in arraylist.
function count_prime_and_NonPrime(n, v) {
    let p = 0, np = 0;
    v.push(new pair(p, np));
    for (let i = 1; i <= n; i++) {
        if (isPrime(i)) {
            p++;
        }
        else {
            np++;
        }
        v.push(new pair(p, np));
    }
    return v;
}

// Function to find number of ways
function NoOfWays(n, a) {
    let v = [];
    let mx = 0;
    for (let i = 0; i < n; i++) {
        mx = Math.max(mx, a[i]);
    }
    v = count_prime_and_NonPrime(mx, v);
    let ans = 1;
    for (let j = 0; j < n; j++) {
        let prime = v[a[j]].first;
        let nonPrime = v[a[j]].second;
        if (isPrime(j + 1)) {
            ans *= nonPrime;
        }
        else {
            ans *= prime;
        }
    }
    return ans;
}

let N = 5;
let A = [2, 3, 4, 8, 5];

// Function call
console.log(NoOfWays(N, A));

// This code is contributed by Saurabh.

Output

16

Time Complexity: O(N * sqrt(M)) where M is the largest element of array
Auxiliary Space: O(N)