February 2022

C2 - Maximum Joy (Hard Version)

Solution

\[\]

Implementation in C++

#include <bits/stdc++.h>

#define ll long long

using namespace std;

int main() {
    int t, n, i;
    ll a[200001], prefix[200001], suffix[200001], sum, rsum, lsum, maxm;
    cin >> t;
    
    while(t--) {
        cin >> n;
        for(i = 0; i < n; ++i)
            cin >> a[i];
        
        //Calculate prefix sum of array a
        sum = 0;
        for(i = 0; i < n; ++i) {
            sum += a[i];
            prefix[i] = sum;
        }
        
        //Calculate suffix sum of array a
        sum = 0, rsum = 0;
        for(i = n-1; i >= 0; --i) {
            sum += a[i];
            suffix[i] = sum;
        }
        
        
        // Calculate the sum of costs for first position prehand
        for(i = 0; i < n; ++i)
           rsum += a[i]*(i+1);
           
           
        // Calculating two sums one from position to all elements 
        // to the right and another sum from start index till the 
        // position of element and finding the maximum among them
        lsum = 0, maxm = 0;
        for(i = 0; i < n; ++i) {
            if(maxm < lsum + rsum)
                maxm = lsum + rsum;
            rsum -= suffix[i];
            if(i != 0) {
                lsum -= prefix[i-1];
                lsum += a[i]*n;
            }
            else
                lsum += a[i]*n;
        }
        
        cout << maxm << "\n";
    }
    return 0;
}
\[\]

Implementation in Java

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

public class Solution {

    public static void main(String[] args) {
        Scanner scan = new Scanner(System.in);
        
        int t, n, i;
        long a[] = new long[200001];
        long prefix[] = new long[200001];
        long suffix[] = new long[200001];
        long sum, rsum, lsum, maxm;
        
        t = scan.nextInt();

        while (t-- > 0) {
            n = scan.nextInt();
            for(i = 0; i < n; ++i)
                a[i] = scan.nextLong();

            //Calculate prefix sum of array a
            sum = 0;
            for(i = 0; i < n; ++i) {
                sum += a[i];
                prefix[i] = sum;
            }

            //Calculate suffix sum of array a
            sum = 0; rsum = 0;
            for(i = n-1; i >= 0; --i) {
                sum += a[i];
                suffix[i] = sum;
            }


            // Calculate the sum of costs for first position prehand
            for(i = 0; i < n; ++i)
               rsum += a[i]*(i+1);


            // Calculating two sums one from position to all elements 
            // to the right and another sum from start index till the 
            // position of element and finding the maximum among them
            lsum = 0; maxm = 0;
            for(i = 0; i < n; ++i) {
                if(maxm < lsum + rsum)
                    maxm = lsum + rsum;
                rsum -= suffix[i];
                if(i != 0) {
                    lsum -= prefix[i-1];
                    lsum += a[i]*n;
                }
                else
                    lsum += a[i]*n;
            }

            System.out.println(maxm);
        }
    }
}
\[\]

Implementation in Python

t = int(input())
for _ in range(t):
    n = int(input())
    a = list(map(int, input().split()))
    
    prefix = []
    suffix = []
    
    for i in range(0, n):
        prefix.append(a[i])
        suffix.append(a[i])
        
    # Calculate prefix sum of array a
    for i in range(1, n):
        prefix[i] += prefix[i-1]
        
    # Calculate suffix sum of array a
    for i in range(n-2, -1, -1):
        suffix[i] += suffix[i+1]
        
    rsum = 0
    
    # Calculate the sum of costs for first position prehand
    for i in range(0, n):
       rsum += a[i]*(i+1)

    # Calculating two sums one from position to all elements 
    # to the right and another sum from start index till the 
    # position of element and finding the maximum among them
    lsum = 0
    maxm = 0
    for i in range(0, n):
        if maxm < lsum + rsum:
            maxm = lsum + rsum
            
        rsum -= suffix[i]
        
        if i != 0:
            lsum -= prefix[i-1]
            lsum += a[i]*n
        else:
            lsum += a[i]*n

    print(maxm)
\[\]

Contest Material