PARTSCORE - Editorial

PROBLEM LINK:

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

Author: raysh07
Tester: sushil2006
Editorialist: iceknight1093

DIFFICULTY:

Simple

PREREQUISITES:

Sorting

PROBLEM:

You’re given an array A and an integer K.
Partition the array into subsets S_1 and S_2, both with size at least K.
Find the maximum possible value of \min(S_1) + \min(S_2) + \max(S_1) + \max(S_2).

EXPLANATION:

Note that no matter how we partition elements into two subarrays, the minimum and maximum elements of A will definitely be present as two of the four elements - they will surely be the minimum/maximum element in whichever subset they’re a part of.

Our aim is hence to maximize the other two elements that show up.
We have one maximum and one minimum to work with.
Of course, the optimal choice for the second maximum is simply the second maximum of A - we only need to decide the second minimum now.

Any element that’s a minimum should have at least K-1 elements greater than it within its own subset, since each subset must have size \geq K.
This means that the absolute largest element we can choose is the K-th maximum element of A.

However, we’ve fixed the two largest elements of A to be in separate subsets, since they each need to be the maximum of their respective subsets.
So, rather than the K-th maximum, we can only use the (K+1)-th maximum element of A.

Thus, to put it simply, the answer is the sum of the following four values:

  1. The minimum element of A.
  2. The maximum element of A.
  3. The second maximum element of A.
  4. The (K+1)-th maximum of A.

It’s not hard to see that creating a partition that uses all these values is indeed achievable. For instance,

  • Set S_1 to include the smallest (N-K-1) elements of A, along with its maximum.
    • The minimum and maximum of this subset are the minimum and maximum of A.
  • Set S_2 to include everything else.
    • The minimum of this subset is the (K+1)-th maximum of A, while its maximum is the second maximum of A.

There is one edge case, however: K = 1.
Our earlier solution differentiated between the minimum and maximum elements of each subset, forcing us to choose 4 distinct elements.
However, K = 1 allows a subset of size 1, and the minimum and maximum elements are equal for such a subset.

For K = 1, we thus set the maximum element of A to be its own subset, and then take everything else into a single subset.
This partition has a value of \min(A) + 2\cdot\max(A) + \text{second\_max}(A) which is optimal.


Finding the min/max/second max/(K+1)-th max are easily done once the array is sorted.

TIME COMPLEXITY:

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

CODE:

Editorialist's code (PyPy3)
for _ in range(int(input())):
    n, k = map(int, input().split())
    a = sorted(list(map(int, input().split())))
    
    ans = a[0] + a[-1] + a[-2] + a[-k-1]
    if k == 1: ans = 2*a[-1] + a[0] + a[-2]
    print(ans)