PARTITION ARRAY INTO K SUBARRAYS TO MAXIMIZE BEAUTY | INFOSYS POWER PROGRAMMER OA

Given N array of an elements of array A. You should partition it into exactly non empty sub arrays such that each element belongs to exactly one subarrays. The beauty of the subarray equals X*X where X is the maximum element in the subarray. For example, the beauty of the subarray [2,4,3] equals 4*4=16.

Output the maximum sum of beauties you can achieve, by partitioning the array A into exactly K subarrays.

Example 1-
N=1
K=1
A=[1,1,1]
Output - 1
Explanation - Because we are only having one subarray that is the array itself where max element is 1.

Example 2-
N=3
K=2
A=[1,2,3]
Output - 13
Explanation - Partition it like [1,2],[3] => 2*2 + 3*3=13

Constraints-
1<=N<=10^5
1<=K<=min(10,N)
1<=A[i]<=10^6

------------------------------------------------------------------------------------------------------------------------------------

vector<vector<long long>> dp;
vector<int> G;

long long helper(vector<int> &arr, int curr, int k){
    if(k==1){
        return (long long)G[curr]*G[curr];
    }
    if(dp[curr][k]!=-1){
        return dp[curr][k];
    }
    long long res=INT_MIN;
    long long mx=INT_MIN;
    for(int i=curr;i<=arr.size()-k;i++){
        mx=max(mx,(long long)arr[i]);
        long long beauty=helper(arr,i+1,k-1);
        if(beauty!=INT_MIN){
            beauty+=(long long)(mx*mx);
            res=max(res,beauty);
        }
    }
    return dp[curr][k]=res;
}

int GetMaxBeauty(int N, int K, vector<int>& arr){
    G=arr;
    for(int i=N-2;i>=0;i--){
        G[i]=max(G[i],G[i+1]);
    }
    dp.resize(N+1,vector<long long>(K+1,-1));
    return helper(arr,0,K);
}  // TLE (10/11) PASSED


---------------------------------------------------------------------------------------------------------------------------------------

JUST PICK K LARGEST ELEMENTS, CUZ WE CAN ALWAYS PARTITION SUCH THAT THESE K ELEMENTS END UP IN K PARTITIONS

int GetMaxBeauty(int N, int K, vector<int>& arr){
    priority_queue<int,vector<int>,greater<int>> pq;
    for(int i=0;i<N;i++){
        pq.push(arr[i]);
        if(pq.size()>K){
            pq.pop();
        }
    }
    int res=0;
    while(!pq.empty()){
        res+=(pq.top()*pq.top());
        pq.pop();
    }
    return res;
}  // ACCEPTED