Maximizing the sum of subsequence of size at most M

Given array A[] of size N, integer M and integer K. The task is to choose at most M elements from the array A[] to form a new array such that if the previous element was chosen is at index i that is A[i] and the current element at index j that is A[j] append A[j] – K * (j – i) to new array. Find the maximum possible sum of elements of the new array.

Note: If no element is chosen before then value of i will be 0 that is A[j] – K * j and it is 1-based indexing.

Examples:

Input: A[] = {3, 2, 5, 4, 6}, M = 2, K = 2
Output: 2
Explanation: It will be optimal to include first and third element of A[] for new array. initially newArray = {} is empty.

• picking first element for new array, previously no element was picked so i = 0 and j = 1 (since we are picking first element) appending A[j] – K * (j – i) = 3 – 2 * (1 – 0) = 1 in new array it becomes newArray = {1}
• picking third element for new array previously we picked element at index i = 1 now we are picking element at index j = 3 so appending A[j] – K * (j – i) = 5 – 2 * (3 – 1) = 5 – 4 = 1 in new array it becomes newArray = {1, 1}
  So the sum of newArray will be 2.

Input: A[] = {1, 1, 1, 1}, M = 3, K = 2
Output: 0
Explanation: It will be optimal not to pick any element in this case so sum is zero.

Approach: To solve the problem follow the below idea:

If we include elements i₁, i₂ ,…, i_k from the array A[] the sum will be:

(A[i₁] − K * i₁) + (A[i₂] − K * (i₂ − i₁)) + … + (A[i_k] − K * (i_k−i_{k – 1})) = (A[i₁] + A[i₂] +…+ A[i_k]) − K * i_k, by maintaining sum of top M non negative elements using priority queue this problem can be solved.

Below are the steps for the above approach:

• Declare priority queue named ms for tracking first M maximum elements from array.
• Declare currentSum = 0 and ans = 0 for tracking current sum during iterating and final answer of the problem.
• Iterate over N elements
  • if A[i] is negative skip the iteration.
  • insert A[i] in priority queue
  • add A[i] in currentSum
  • if size of priority queue is greater than M than delete smallest element from priority queue and subtract that element from currentSum
  • Update the answer in ans variable
• return final answer ans.

Below is the implementation of the above approach:

C++

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

// Function to maximize the sum of subsequence
// of size at most M
int maximizeSum(int A[], int N, int M, int K)
{

    // Declaring priority queue for tracking
    // sum of at most M elements of array
    multiset<int> ms;

    // Declaring currentSum variable for tracking
    // sum and ans variable for tracking
    // final answer
    int currentSum = 0, ans = 0;

    // Iterating over array
    for (int i = 0; i < N; i++) {

        // Skip the iteration if number
        // is negative
        if (A[i] < 0)
            continue;

        // Inserting element in priority queue
        ms.insert(A[i]);

        // Adding A[i] in current sum
        currentSum = currentSum + A[i];

        // If size of priority queue exceeds
        // M delete smallest element
        if (ms.size() > M) {

            // Subtracting smallest element
            // from current sum
            currentSum = currentSum - *ms.begin();

            // Erasing smallest element from
            // priority queue
            ms.erase(ms.begin());
        }

        // Updating the answer
        ans = max(ans, currentSum - (i + 1) * K);
    }

    // Returning the final answer
    return ans;
}

// Driver Code
int32_t main()
{

    // Input 1
    int M = 2, K = 2;
    int A[] = { 3, 2, 5, 4, 6 }, N = 5;

    // Function Call
    cout << maximizeSum(A, N, M, K) << endl;

    // Input 2
    int M1 = 3, K1 = 2;
    int A1[] = { 1, 1, 1, 1 }, N1 = 4;

    // Function Call
    cout << maximizeSum(A1, N1, M1, K1) << endl;

    return 0;
}

Java

import java.util.PriorityQueue;

public class Main {
    public static long maximizeSum(int[] A, int N, int M, int K) {
        // Declaring priority queue for tracking
        // sum of at most M elements of array
        PriorityQueue<Integer> pq = new PriorityQueue<>();

        // Declaring currentSum variable for tracking
        // sum and ans variable for tracking
        // final answer
        long currentSum = 0, ans = 0;

        // Iterating over the array
        for (int i = 0; i < N; i++) {
            // Skip the iteration if the number is negative
            if (A[i] < 0)
                continue;

            // Inserting the element in the priority queue
            pq.offer(A[i]);

            // Adding A[i] to the current sum
            currentSum += A[i];

            // If the size of the priority queue exceeds M, remove the smallest element
            if (pq.size() > M) {
                // Subtracting the smallest element from the current sum
                currentSum -= pq.poll();
            }

            // Updating the answer
            ans = Math.max(ans, currentSum - (i + 1) * K);
        }

        // Returning the final answer
        return ans;
    }

    public static void main(String[] args) {
        // Input 1
        int M = 2, K = 2;
        int[] A = {3, 2, 5, 4, 6};
        int N = 5;

        // Function Call
        System.out.println(maximizeSum(A, N, M, K));

        // Input 2
        int M1 = 3, K1 = 2;
        int[] A1 = {1, 1, 1, 1};
        int N1 = 4;

        // Function Call
        System.out.println(maximizeSum(A1, N1, M1, K1));
    }
}
//Contributed by Aditi Tyagi

Python3

import heapq

def maximize_sum(A, N, M, K):
    # Declaring priority queue for tracking
    # sum of at most M elements of array
    pq = []

    # Declaring currentSum variable for tracking
    # sum and ans variable for tracking
    # final answer
    current_sum = 0
    ans = 0

    # Iterating over the array
    for i in range(N):
        # Skip the iteration if the number is negative
        if A[i] < 0:
            continue

        # Inserting the element in the priority queue
        heapq.heappush(pq, A[i])

        # Adding A[i] to the current sum
        current_sum += A[i]

        # If the size of the priority queue exceeds M, remove the smallest element
        if len(pq) > M:
            # Subtracting the smallest element from the current sum
            current_sum -= heapq.heappop(pq)

        # Updating the answer
        ans = max(ans, current_sum - (i + 1) * K)

    # Returning the final answer
    return ans

# Input 1
M = 2
K = 2
A = [3, 2, 5, 4, 6]
N = 5

# Function Call
print(maximize_sum(A, N, M, K))

# Input 2
M1 = 3
K1 = 2
A1 = [1, 1, 1, 1]
N1 = 4

# Function Call
print(maximize_sum(A1, N1, M1, K1))
#This code is contributed by rohit singh

C#

using System;
using System.Collections.Generic;

public class Program
{
      // Function to maximize the sum of subsequence
    // of size at most M
    public static int MaximizeSum(int[] A, int N, int M, int K)
    {
        SortedSet<int> ss = new SortedSet<int>();

          // Declaring currentSum variable for tracking
        // sum and ans variable for tracking
        // final answer
        int currentSum = 0, ans = 0;

          // Iterating over array
        for (int i = 0; i < N; i++)
        {
              // Skip the iteration if number
            // is negative
            if (A[i] < 0)
                continue;

              // Inserting element in priority queue
            ss.Add(A[i]);

              // Adding A[i] in current sum
            currentSum += A[i];

              // If size of priority queue exceeds
            // M delete smallest element
            if (ss.Count > M)
            {
                  // Subtracting smallest element
                // from current sum
                currentSum -= ss.Min;

                  // Erasing smallest element from
                // priority queue
                ss.Remove(ss.Min);
            }

              // Updating the answer
            ans = Math.Max(ans, currentSum - (i + 1) * K);
        }

          // Returning the final answer
        return ans;
    }

      // Driver Code
    public static void Main(string[] args)
    {
          // Input 1
        int M = 2, K = 2;
        int[] A = { 3, 2, 5, 4, 6 };
        int N = 5;
        Console.WriteLine(MaximizeSum(A, N, M, K));

          // Input 2
        int M1 = 3, K1 = 2;
        int[] A1 = { 1, 1, 1, 1 };
        int N1 = 4;
        Console.WriteLine(MaximizeSum(A1, N1, M1, K1));
    }
}

Javascript

// Function to maximize the sum of subsequence
// of size at most M
function maximizeSum(A, N, M, K) {

    // Declaring priority queue for tracking
    // sum of at most M elements of array
    let ms = new Set();

    // Declaring currentSum variable for tracking
    // sum and ans variable for tracking
    // final answer
    let currentSum = 0;
    let ans = 0;

    // Iterating over array
    for (let i = 0; i < N; i++) {

        // Skip the iteration if number
        // is negative
        if (A[i] < 0)
            continue;

        // Inserting element in priority queue
        ms.add(A[i]);

        // Adding A[i] in current sum
        currentSum = currentSum + A[i];

        // If size of priority queue exceeds
        // M delete smallest element
        if (ms.size > M) {

            // Subtracting smallest element
            // from current sum
            currentSum = currentSum - Math.min(...ms);

            // Erasing smallest element from
            // priority queue
            ms.delete(Math.min(...ms));
        }

        // Updating the answer
        ans = Math.max(ans, currentSum - (i + 1) * K);
    }

    // Returning the final answer
    return ans;
}

// Driver Code
// Input 1
let M = 2, K = 2;
let A = [3, 2, 5, 4, 6];
let N = 5;

// Function Call
console.log(maximizeSum(A, N, M, K));

// Input 2
let M1 = 3, K1 = 2;
let A1 = [1, 1, 1, 1];
let N1 = 4;

// Function Call
console.log(maximizeSum(A1, N1, M1, K1));

Output

2
0

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