Apply Operations to Maximize Score | Leetcode 2818

class Solution {
    #define ll long long
    #define pii pair<ll,ll>
    ll MOD = 1e9+7;

    void calculateScore(vector<int>& nums,vector<ll>& score){
        for(ll ele: nums){
            ll count=0;
            for(ll i=2;i*i<ele;++i){
                if(ele%i==0)
                    count++;
                while(ele%i==0)
                    ele/=i;
            }
            if(ele>1)   count++;
            score.push_back(count);
        }
    }
    void calculateSubarrayCountPerScore(vector<ll>& score,vector<ll>& subarray_count){
        //Calculate Previous Greater or Equal Element
        vector<ll> pge;
        stack<ll> st;
        ll n=score.size();
        for(ll i=0;i<n;++i){
            while(!st.empty() and score[st.top()]<score[i])
                st.pop();
            if(st.empty())  pge.push_back(-1);
            else            pge.push_back(st.top());
            st.push(i);
        }

        //Parsse from right (next greater element) and count subarrays
        st = stack<ll>();
        ll count;
        for(ll i=n-1;i>=0;--i){
            count=0;
            while(!st.empty() and score[st.top()]<=score[i])
                st.pop();
            if(st.empty())  count = (n-i)*(i-pge[i]);
            else            count = (st.top()-i)*(i-pge[i]);
            
            st.push(i);
            subarray_count[i]=count;
        }
    }
    ll binaryExponentiation(long long a,long long b){
        //Find pow(a,b) mod (10^9+7)
        ll res=1;
        while(b){
            if(b&1)
                res = (res * a) % MOD;
            a = (a * a) % MOD;
            b /= 2;
        }
        return res % MOD;
    }
public:
    int maximumScore(vector<int>& nums, ll k) {
        ll n=nums.size();
        vector<ll> score;
        calculateScore(nums,score);

        vector<ll> subarray_count(n);
        calculateSubarrayCountPerScore(score,subarray_count);
        
        //Push (score,subarray) pair in maxheap
        priority_queue<pii> maxheap;
        for(ll i=0;i<n;++i)
            maxheap.push(make_pair(nums[i],i));
        
        long long res=1;
        while(k>0){
            pii curr = maxheap.top();
            maxheap.pop();
            res = (res * binaryExponentiation(curr.first,min(k,subarray_count[curr.second]))) % MOD;
            k -= subarray_count[curr.second];
        }
        return res;
    }
};
/*
//JAVA
import java.util.*;

class Solution {
    private static final long MOD = 1000000007;

    private void calculateScore(int[] nums, List<Long> score) {
        for (int ele : nums) {
            long count = 0;
            for (int i = 2; i * i <= ele; ++i) {
                if (ele % i == 0) {
                    count++;
                    while (ele % i == 0) {
                        ele /= i;
                    }
                }
            }
            if (ele > 1) count++;
            score.add(count);
        }
    }

    private void calculateSubarrayCountPerScore(List<Long> score, long[] subarrayCount) {
        int n = score.size();
        int[] pge = new int[n];
        Arrays.fill(pge, -1);
        Deque<Integer> stack = new ArrayDeque<>();

        // Calculate Previous Greater Element (PGE)
        for (int i = 0; i < n; ++i) {
            while (!stack.isEmpty() && score.get(stack.peek()) < score.get(i)) {
                stack.pop();
            }
            if (!stack.isEmpty()) {
                pge[i] = stack.peek();
            }
            stack.push(i);
        }

        // Calculate Next Greater Element (NGE) and subarray count
        stack.clear();
        for (int i = n - 1; i >= 0; --i) {
            while (!stack.isEmpty() && score.get(stack.peek()) <= score.get(i)) {
                stack.pop();
            }
            int nge = stack.isEmpty() ? n : stack.peek();
            long count = (long) (nge - i) * (i - pge[i]);
            subarrayCount[i] = count % MOD;
            stack.push(i);
        }
    }

    private long binaryExponentiation(long a, long b) {
        long res = 1;
        a %= MOD;
        while (b > 0) {
            if ((b & 1) == 1) {
                res = (res * a) % MOD;
            }
            a = (a * a) % MOD;
            b >>= 1;
        }
        return res;
    }

    public int maximumScore(int[] nums, int k) {
        int n = nums.length;
        List<Long> score = new ArrayList<>();
        calculateScore(nums, score);

        long[] subarrayCount = new long[n];
        calculateSubarrayCountPerScore(score, subarrayCount);

        // Use a max-heap to prioritize the largest numbers
        PriorityQueue<int[]> maxHeap = new PriorityQueue<>((a, b) -> b[0] - a[0]);
        for (int i = 0; i < n; ++i) {
            maxHeap.offer(new int[]{nums[i], i});
        }

        long res = 1;
        while (k > 0 && !maxHeap.isEmpty()) {
            int[] curr = maxHeap.poll();
            int num = curr[0];
            int idx = curr[1];
            long cnt = Math.min(k, subarrayCount[idx]);
            res = (res * binaryExponentiation(num, cnt)) % MOD;
            k -= cnt;
        }
        return (int) res;
    }
}

#Python
import heapq
from collections import deque

MOD = 10**9 + 7

class Solution:
    def calculateScore(self, nums):
        score = []
        for ele in nums:
            count = 0
            i = 2
            while i * i <= ele:
                if ele % i == 0:
                    count += 1
                    while ele % i == 0:
                        ele //= i
                i += 1
            if ele > 1:
                count += 1
            score.append(count)
        return score

    def calculateSubarrayCountPerScore(self, score):
        n = len(score)
        pge = [-1] * n
        stack = deque()

        # Calculate Previous Greater Element (PGE)
        for i in range(n):
            while stack and score[stack[-1]] < score[i]:
                stack.pop()
            if stack:
                pge[i] = stack[-1]
            stack.append(i)

        # Calculate Next Greater Element (NGE) and subarray count
        stack.clear()
        subarray_count = [0] * n
        for i in range(n - 1, -1, -1):
            while stack and score[stack[-1]] <= score[i]:
                stack.pop()
            nge = stack[-1] if stack else n
            count = (nge - i) * (i - pge[i])
            subarray_count[i] = count % MOD
            stack.append(i)
        return subarray_count

    def binaryExponentiation(self, a, b):
        res = 1
        a %= MOD
        while b > 0:
            if b & 1:
                res = (res * a) % MOD
            a = (a * a) % MOD
            b >>= 1
        return res

    def maximumScore(self, nums, k):
        n = len(nums)
        score = self.calculateScore(nums)
        subarray_count = self.calculateSubarrayCountPerScore(score)

        # Use a max-heap to prioritize the largest numbers
        max_heap = [(-nums[i], i) for i in range(n)]
        heapq.heapify(max_heap)

        res = 1
        while k > 0 and max_heap:
            num, idx = heapq.heappop(max_heap)
            num = -num
            cnt = min(k, subarray_count[idx])
            res = (res * self.binaryExponentiation(num, cnt)) % MOD
            k -= cnt
        return res

*/

//Your code was not working
// Working Java Code
import java.util.*;

class Solution {
private static final long MOD = 1000000007;

private void calculateScore(int[] nums, List<Long> score) {
    for (int ele : nums) {
        long count = 0;
        int num = ele;
        for (int i = 2; i * i <= num; ++i) {
            if (num % i == 0) {
                count++;
                while (num % i == 0) {
                    num /= i;
                }
            }
        }
        if (num > 1) count++;
        score.add(count);
    }
}

private void calculateSubarrayCountPerScore(List<Long> score, long[] subarrayCount) {
    int n = score.size();
    int[] pge = new int[n];
    Arrays.fill(pge, -1);
    Deque<Integer> stack = new ArrayDeque<>();

    for (int i = 0; i < n; ++i) {
        while (!stack.isEmpty() && score.get(stack.peek()) < score.get(i)) {
            stack.pop();
        }
        if (!stack.isEmpty()) {
            pge[i] = stack.peek();
        }
        stack.push(i);
    }

    stack.clear();
    for (int i = n - 1; i >= 0; --i) {
        while (!stack.isEmpty() && score.get(stack.peek()) <= score.get(i)) {
            stack.pop();
        }
        int nge = stack.isEmpty() ? n : stack.peek();
        long count = (long) (nge - i) * (i - pge[i]);
        subarrayCount[i] = count;
        stack.push(i);
    }
}

private long binaryExponentiation(long a, long b) {
    long res = 1;
    a %= MOD;
    while (b > 0) {
        if ((b & 1) == 1) {
            res = (res * a) % MOD;
        }
        a = (a * a) % MOD;
        b >>= 1;
    }
    return res;
}

public int maximumScore(List<Integer> nums, int k) {
    int n = nums.size();
    List<Long> score = new ArrayList<>();
    int[] numArray = nums.stream().mapToInt(i -> i).toArray();
    calculateScore(numArray, score);

    long[] subarrayCount = new long[n];
    calculateSubarrayCountPerScore(score, subarrayCount);

    PriorityQueue<int[]> maxHeap = new PriorityQueue<>((a, b) -> Integer.compare(b[0], a[0]));
    for (int i = 0; i < n; ++i) {
        maxHeap.offer(new int[]{numArray[i], i});
    }

    long res = 1;
    while (k > 0 && !maxHeap.isEmpty()) {
        int[] curr = maxHeap.poll();
        int num = curr[0];
        int idx = curr[1];
        long cnt = Math.min(k, subarrayCount[idx]);
        res = (res * binaryExponentiation(num, cnt)) % MOD;
        k -= cnt;
        if (k == 0) break;
    }
    return (int) res;
}

}