Count subarrays having product equal to the power of a given Prime Number

Given an array arr[] of size N and an integer M, the task is to count the number of subarrays having product of its elements equal to the power of M, where M is a prime number.

Examples:

Input: arr[] = {2, 2, 2, 2}, M = 2
Output: 10
Explanation: All possible non-empty subarrays having product equal to the power of M = (4 * (4 + 1)) / 2 = 10

Input: arr[] = {1, 1, 1, 3}, M = 3
Output: 10

Naive Approach: The simplest approach is to generate all possible subarrays of the array arr[] and for each subarray, check if their product is a power of M or not. If found to be true, then increment the count of such subarrays by 1. Finally, print the count obtained. 

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

Efficient Approach: The optimal idea is based upon the fact that for the product of any subarray to be equal to the power of M, all the elements in the subarray must also a power of M, since M is a prime number. Follow the steps below to solve the given problem:

• Initialize variable, say ans, to store the count of required subarrays
• Initialize a variable, say cnt, to store the count of consecutive numbers that are a power of M.
• Traverse the array over the range of indices [0, N – 1] using a variable, say i, and perform the following operations:
  • If arr[i] is a power of M. Increment cnt by 1. Update ans = ans + (cnt * (cnt – 1)) / 2
  • Otherwise, update cnt = 0.
• After completing the above steps, print the value of cnt.

Below is the implementation of the above approach:

C++

// C++ program for the above approach
#include <bits/stdc++.h>
using namespace std;

// Function to check if y
// is a power of m or not
bool isPower(int m, int y)
{
    // Calculate log y base m and store
    // it in a variable with integer datatype
    int res1 = log(y) / log(m);

    // Calculate log y base m and store
    // it in a variable with double datatype
    double res2 = log(y) / log(m);

    // If res1 and res2 are equal, return
    // True. Otherwise, return false
    return (res1 == res2);
}

// Function to count the number of subarrays
// having product of elements equal to a
// power of m, where m is a prime number
int numSub(int arr[], int n, int m)
{
    // Stores the count of
    // subarrays required
    int ans = 0;

    // Stores current sequence of
    // consecutive array elements
    // which are a multiple of m
    int cnt = 0;

    // Traverse the array
    for (int i = 0; i < n; i++) {

        // If arr[i] is a power of M
        if (isPower(m, arr[i])) {

            // Increment cnt
            cnt++;

            // Update ans
            ans += (cnt * (cnt - 1)) / 2;
        }
        else {

            // Update cnt
            cnt = 0;
        }
    }

    // Return the count
    // of subarrays
    return ans;
}

// Driver Code
int main()
{
    // Input
    int arr[] = { 1, 1, 1, 3 };
    int m = 3;
    int n = sizeof(arr) / sizeof(arr[0]);

    cout << numSub(arr, n, m);

    return 0;
}

Java

// Java code of above approach
import java.util.*;
class GFG
{

// Function to check if y
// is a power of m or not
static boolean isPower(int m, int y)
{

    // Calculate log y base m and store
    // it in a variable with integer datatype
    int res1 = (int)Math.log(y) / (int)Math.log(m);

    // Calculate log y base m and store
    // it in a variable with double datatype
    double res2 = (int)Math.log(y) / (int)Math.log(m);

    // If res1 and res2 are equal, return
    // True. Otherwise, return false
    return (res1 == res2);
}

// Function to count the number of subarrays
// having product of elements equal to a
// power of m, where m is a prime number
static int numSub(int arr[], int n, int m)
{

    // Stores the count of
    // subarrays required
    int ans = 0;

    // Stores current sequence of
    // consecutive array elements
    // which are a multiple of m
    int cnt = 0;

    // Traverse the array
    for (int i = 0; i < n; i++) {

        // If arr[i] is a power of M
        if (isPower(m, arr[i])) {

            // Increment cnt
            cnt++;

            // Update ans
            ans += (cnt * (cnt - 1)) / 2;
        }
        else {

            // Update cnt
            cnt = 0;
        }
    }

    // Return the count
    // of subarrays
    return ans;
}

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

     // Input
    int arr[] = { 1, 1, 1, 3 };
    int m = 3;
    int n = arr.length;

   System.out.println(numSub(arr, n, m));
    }
}

// This code is contributed by offbeat

Python3

# Python 3 program for the above approach
from math import log

# Function to check if y
# is a power of m or not
def isPower(m, y):

    # Calculate log y base m and store
    # it in a variable with integer datatype
    res1 = log(y) // log(m)

    # Calculate log y base m and store
    # it in a variable with double datatype
    res2 = log(y) // log(m)

    # If res1 and res2 are equal, return
    # True. Otherwise, return false
    return (res1 == res2)

# Function to count the number of subarrays
# having product of elements equal to a
# power of m, where m is a prime number
def numSub(arr, n, m):

    # Stores the count of
    # subarrays required
    ans = 0

    # Stores current sequence of
    # consecutive array elements
    # which are a multiple of m
    cnt = 0

    # Traverse the array
    for i in range(n):

        # If arr[i] is a power of M
        if (isPower(m, arr[i])):

            # Increment cnt
            cnt += 1

            # Update ans
            ans += (cnt * (cnt - 1)) // 2

        else:
            # Update cnt
            cnt = 0

    # Return the count
    # of subarrays
    return ans

# Driver Code
if __name__ == '__main__':

    # Input
    arr =  [1, 1, 1, 3]
    m = 3
    n = len(arr)
    print(numSub(arr, n, m))

    # This code is contributed by bgangwar59.

C#

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

class GFG{

// Function to check if y
// is a power of m or not
static bool isPower(int m, int y)
{

    // Calculate log y base m and store
    // it in a variable with integer datatype
    int res1 = (int)Math.Log(y) / (int)Math.Log(m);

    // Calculate log y base m and store
    // it in a variable with double datatype
    double res2 = (int)Math.Log(y) / (int)Math.Log(m);

    // If res1 and res2 are equal, return
    // True. Otherwise, return false
    return (res1 == res2);
}

// Function to count the number of subarrays
// having product of elements equal to a
// power of m, where m is a prime number
static int numSub(int[] arr, int n, int m)
{

    // Stores the count of
    // subarrays required
    int ans = 0;

    // Stores current sequence of
    // consecutive array elements
    // which are a multiple of m
    int cnt = 0;

    // Traverse the array
    for(int i = 0; i < n; i++)
    {

        // If arr[i] is a power of M
        if (isPower(m, arr[i]))
        {

            // Increment cnt
            cnt++;

            // Update ans
            ans += (cnt * (cnt - 1)) / 2;
        }
        else
        {

            // Update cnt
            cnt = 0;
        }
    }

    // Return the count
    // of subarrays
    return ans;
}

// Driver code
public static void Main()
{

    // Input
    int[] arr = { 1, 1, 1, 3 };
    int m = 3;
    int n = arr.Length;

    Console.Write(numSub(arr, n, m));
}
}

// This code is contributed by subhammahato348

Javascript

<script>
// Javascript program for the above approach

// Function to check if y
// is a power of m or not
function isPower(m, y)
{

    // Calculate log y base m and store
    // it in a variable with integer datatype
    let res1 = parseInt(Math.log(y) / Math.log(m));

    // Calculate log y base m and store
    // it in a variable with double datatype
    let res2 = Math.log(y) / Math.log(m);

    // If res1 and res2 are equal, return
    // True. Otherwise, return false
    return (res1 == res2);
}

// Function to count the number of subarrays
// having product of elements equal to a
// power of m, where m is a prime number
function numSub(arr, n, m)
{
    // Stores the count of
    // subarrays required
    let ans = 0;

    // Stores current sequence of
    // consecutive array elements
    // which are a multiple of m
    let cnt = 0;

    // Traverse the array
    for (let i = 0; i < n; i++) {

        // If arr[i] is a power of M
        if (isPower(m, arr[i])) {

            // Increment cnt
            cnt++;

            // Update ans
            ans += parseInt((cnt * (cnt - 1)) / 2);
        }
        else {

            // Update cnt
            cnt = 0;
        }
    }

    // Return the count
    // of subarrays
    return ans;
}

// Driver Code
    // Input
    let arr = [ 1, 1, 1, 3 ];
    let m = 3;
    let n = arr.length;

    document.write(numSub(arr, n, m));

  // This code is contributed by subhammahato348.
</script>

Output:

10

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