EXSHAKE - Editorial

PROBLEM LINK:

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

Author: am_aadvik
Tester: iceknight1093
Editorialist: iceknight1093

DIFFICULTY:

TBD

PREREQUISITES:

Sorting

PROBLEM:

N groups of people arrive at a symposium, group i has A_i people.
Each group can be either invited in or left out.
If P people in total enter, there will be P\cdot (P-1)/2 handshakes.

Find the maximum number of groups who can be invited in while keeping the number of handshakes to at most M.

EXPLANATION:

If we decide to invite k groups of people in, then in order to minimize the number of handshakes it’s clearly optimal to invite the smallest k groups of people.

This applies to any choice of k, and so we can simply try all of them.

Sort the array A.
Then, for each k = 1, 2, 3, \ldots, N, the number of people we invite (assuming k groups) will be exactly

P_k = A_1 + A_2 + \ldots + A_k

With P_0 = 0 since we can choose to not invite any groups.

Now, simply find the largest k such that P_k \cdot (P_k-1) / 2 \le M and that k is the answer.

TIME COMPLEXITY:

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

CODE:

Editorialist's code (PyPy3)
for _ in range(int(input())):
    n, m = map(int, input().split())
    a = list(map(int,input().split()))
    a.sort()
    
    p = 0
    for i in range(n):
        p += a[i]
        if p*(p-1)//2 > m:
            print(i)
            break
    else:
        print(n)