REDBLUE - Editorial

editorial
, , ,
1

PROBLEM LINK:

Practice
Contest: Division 1
Contest: Division 2
Contest: Division 3
Contest: Division 4

Author: raysh07
Tester: iceknight1093
Editorialist: iceknight1093

DIFFICULTY:

Simple

PREREQUISITES:

Sorting

PROBLEM:

You have an array with N elements. Each must be colored either red or blue.
Let C_R, C_B denote the counts of red and blue elements, and S_R, S_B denote their sums.
Maximize C_R\cdot S_B + C_B\cdot S_R.

EXPLANATION:

Let T = \sum_{i=1}^N A_i denote the total sum of all elements of the array.

Every element must be colored either red or blue.
So, we have

  • C_R + C_B = N
  • S_R + S_B = T

We can use this property to rewrite C_R\cdot S_B + C_B\cdot S_R in terms of the constants N and T, turning it into

C_R\cdot (T - S_R) + (N-C_R)\cdot S_R = C_R\cdot T + N\cdot S_R - 2\cdot C_R\cdot S_R

Now, the value of C_R can only be an integer in the range [0, N], since it represents the number of red elements.

So, suppose C_R is fixed.
What should S_R be?
Looking at the equation, we can group the terms to obtain

S_R\cdot (N - 2C_R) + C_R\cdot T

This is a linear equation in S_R, meaning if we plot its value with respect to S_R the resulting graph will look like a line.

We’re aiming to maximize S_R, and for a line, that happens when S_R is either as small as possible, or as large as possible (the exact one depends on the sign of N - 2C_R but there’s no real need to do that work, as we’ll see later.)

Making S_R as small as possible means choosing the smallest C_R elements to paint red; while making making it as large as possible means choosing the largest C_R elements to become red.

Since red and blue are symmetric, in the end all we’re doing is giving some (sorted) prefix of elements one color, and all other elements the remaining color.

So, we can simply try all choices for this split and take the best one!

To summarize the above, the solution is as follows.

First, sort the array, so that A_1\le\ldots\le A_N.

Then, for each i = 1, 2, \ldots, N, give the first i elements one color and the remaining elements the other color.
This makes the overall cost equal to

i\cdot (A_{i+1} + \ldots + A_N) + (N-i)\cdot (A_1 + \ldots + A_i)

Take the best of this across all i.
By maintaining a running sum of the A_i values, this can be done in linear time after sorting.

TIME COMPLEXITY:

\mathcal{O}(N\log N) per testcase.

CODE:

Editorialist's code (PyPy3) 
for _ in range(int(input())):
    n = int(input())
    a = list(map(int, input().split()))

    a.sort()
    s = sum(a)
    ans, cur = 0, 0
    for i in range(n):
        ans = max(ans, i*(s-cur) + (n-i)*cur)
        cur += a[i]
    print(ans)