Problem:


Score : 500 points

Problem Statement

Takahashi has come to a party as a special guest.
There are N ordinary guests at the party. The i-th ordinary guest has a power of A_i.

Takahashi has decided to perform M handshakes to increase the happiness of the party (let the current happiness be 0).
A handshake will be performed as follows:

Takahashi chooses one (ordinary) guest x for his left hand and another guest y for his right hand (x and y can be the same).

Then, he shakes the left hand of Guest x and the right hand of Guest y simultaneously to increase the happiness by A_x+A_y.

However, Takahashi should not perform the same handshake more than once. Formally, the following condition must hold:

Assume that, in the k-th handshake, Takahashi shakes the left hand of Guest x_k and the right hand of Guest y_k. Then, there is no pair p, q (1 \leq p < q \leq M) such that (x_p,y_p)=(x_q,y_q).

What is the maximum possible happiness after M handshakes?

Constraints

1 \leq N \leq 10^5

1 \leq M \leq N^2

1 \leq A_i \leq 10^5

All values in input are integers.

Input

Input is given from Standard Input in the following format:

N M
A_1 A_2 ... A_N

Output

Print the maximum possible happiness after M handshakes.

Sample Input 1

5 3
10 14 19 34 33

Sample Output 1

202

Let us say that Takahashi performs the following handshakes:

In the first handshake, Takahashi shakes the left hand of Guest 4 and the right hand of Guest 4.

In the second handshake, Takahashi shakes the left hand of Guest 4 and the right hand of Guest 5.

In the third handshake, Takahashi shakes the left hand of Guest 5 and the right hand of Guest 4.

Then, we will have the happiness of (34+34)+(34+33)+(33+34)=202.

We cannot achieve the happiness of 203 or greater, so the answer is 202.

Sample Input 2

9 14
1 3 5 110 24 21 34 5 3

Sample Output 2

1837

Sample Input 3

9 73
67597 52981 5828 66249 75177 64141 40773 79105 16076

Sample Output 3

8128170



Sample input:
5 3
10 14 19 34 33


Sample output:
202

Submission status: Runtime Error

Solution code:
class NTT:
    def __init__(self, convolution_rank, binary_level):
        self.cr = convolution_rank
        self.bl = binary_level
        self.modulo_list = [1]*convolution_rank
        self.primal_root_list = [1]*convolution_rank
        self.bin_list = [1]*(binary_level+1)
        self.primal_base_matrix = [[0 for _ in range(binary_level+1)] for __ in range(convolution_rank)]
        self.inverse_base_matrix = [[0 for _ in range(binary_level+1)] for __ in range(convolution_rank)]
        for i in range(binary_level):
            self.bin_list[i+1] = self.bin_list[i] * 2
    
    def doubling(self, n, m, modulo=0):
        y = 1
        tmp = m
        bas = n
        while tmp:
            if tmp % 2:
                y *= bas
                if modulo:
                    y %= modulo
            bas *= bas
            if modulo:
                bas %= modulo
            tmp >>= 1
        return y
    
    def powlimit(self, n):
        y = 1
        cnt = 0
        while y < n:
            y *= 2
            cnt += 1
        return y, cnt
    
    def IsPrime(self, num):
        p = 2
        if num <= 1:
            return False
        while p * p <= num:
            if num % p == 0:
                return False
            p += 1
        return True
    
    def extgcd(self, a, b, c):
        x, y, u, v, k, l = 1, 0, 0, 1, a, b
        while l:
            x, y, u, v = u, v, x - u * (k // l), y - v * (k // l)
            k, l = l, k % l
        if c % k:
            return "No Solution"
        return x * c, y * c
    
    def inved(self, a, modulo):
        x, y = self.extgcd(a, modulo, 1)
        return (x+modulo)%modulo
    
    def root_manual(self):
        for i in range(self.cr):
            r = 1
            flg = True
            while flg:
                fflg = True
                for j in range(self.bl):
                    if self.doubling(r, self.bin_list[j], self.modulo_list[i]) == 1:
                        fflg = False
                if self.doubling(r, self.bin_list[-1], self.modulo_list[i]) != 1:
                    fflg = False
                if fflg:
                    flg = False
                else:
                    r += 1
                    if r >= self.modulo_list[i]:
                        break
            self.primal_root_list[i] = r
    
    def make_prime_root(self):
        cnt = 0
        j = 1
        last = self.bin_list[-1]
        while cnt < self.cr:
            if self.IsPrime(j*last+1):
                flg = True
                r = 1
                while flg:
                    fflg = True
                    for i in range(self.bl):
                        if self.doubling(r, self.bin_list[i], j*last+1) == 1:
                            fflg = False
                    if self.doubling(r, last, j*last+1) != 1:
                        fflg = False
                    if fflg:
                        flg = False
                    else:
                        r += 1
                        if r >= j*last+1:
                            break
                if flg==False:
                    self.modulo_list[cnt] = j*last+1
                    self.primal_root_list[cnt] = r
                    cnt += 1
            j += 2
    
    def make_basis(self):
        for i in range(self.cr):
            for j in range(self.bl):
                tmp = self.doubling(2, self.bl-j)
                self.primal_base_matrix[i][j] = self.doubling(self.primal_root_list[i], tmp, self.modulo_list[i])
                self.inverse_base_matrix[i][j] = self.inved(self.primal_base_matrix[i][j], self.modulo_list[i])
    
    def NTT(self, f, n, idx, depth, inverse=False, surface=True):
        mod = self.modulo_list[idx]
        i2 = self.inved(2, mod)
        if inverse:
            seed = self.primal_base_matrix[idx][depth]
        else:
            seed = self.inverse_base_matrix[idx][depth]
        if n == 1:
            return f
        if n == 2:
            xf = [(f[0] + f[1])%mod, (f[0] - f[1])%mod]
            if inverse * surface:
                xf[0] *= i2
                xf[0] %= mod
                xf[1] *= i2
                xf[1] %= mod
            return xf
        hn = n // 2
        fe = [f[2*i+0] for i in range(hn)]
        fo = [f[2*i+1] for i in range(hn)]
        fe = self.NTT(fe, hn, idx, depth-1, inverse, False)
        fo = self.NTT(fo, hn, idx, depth-1, inverse, False)
        xf = [0 for _ in range(n)]
        grow = 1
        for i in range(hn):
            right = (fo[i] * grow) % mod
            xf[i+0*hn] = (fe[i] + right) % mod
            xf[i+1*hn] = (fe[i] - right) % mod
            grow *= seed
            grow %= mod
        if inverse * surface:
            invn = self.inved(n, mod)
            for i in range(n):
                xf[i] *= invn
                xf[i] %= mod
        return xf
    
    def second_NTT(self, f, n, idx, depth, inverse=False):
        res = [0 for _ in range(n)]
        tmp = [0 for _ in range(n)]
        MOD = self.modulo_list[idx]
        ipl = self.inved(n, MOD)
        for i in range(n):
            bas = 1
            pos = 0
            for j in range(depth, 0, -1):
                pos += bas * ((i>>(j-1)) % 2)
                bas *= 2
            res[i] = f[pos]
        for i in range(depth):
            grow = 1
            seed = inverse * self.primal_base_matrix[idx][i+1] + (1 - inverse) * self.inverse_base_matrix[idx][i+1]
            for k in range(1<<i):
                for j in range(1<<(depth-i-1)):
                    tmp[j*(1<<(i+1))+k+0*(1<<i)] = (res[j*(1<<(i+1))+k] + grow * res[j*(1<<(i+1))+k+(1<<i)]) % MOD
                    tmp[j*(1<<(i+1))+k+1*(1<<i)] = (res[j*(1<<(i+1))+k] - grow * res[j*(1<<(i+1))+k+(1<<i)]) % MOD
                grow *= seed
                grow %= MOD
            for j in range(n):
                res[j] = tmp[j]
        if inverse:
            for i in range(n):
                res[i] *= ipl
                res[i] %= MOD
        return res
    
    def CRT(self, num, a_list, m_list):
        r = a_list[0]
        bas = m_list[0]
        x, y = 0, 0
        for i in range(1, num):
            x, y = self.extgcd(bas, -m_list[i], a_list[i]-r)
            r += bas * x
            bas *= m_list[i]
        return r % bas

N, M = map(int, input().split())
ntt = NTT(2, 27)
ntt.make_prime_root()
ntt.make_basis()
A = list(map(int, input().split()))
m = max(A)+1
pl, cnt = ntt.powlimit(2*m)
X = [0 for _ in range(pl)]
for i in range(N):
    X[A[i]] += 1
X1 = ntt.second_NTT(X, pl, 0, cnt)
X2 = ntt.second_NTT(X, pl, 1, cnt)
Z1 = [X1[i]*X1[i]%ntt.modulo_list[0] for i in range(pl)]
Z2 = [X2[i]*X2[i]%ntt.modulo_list[1] for i in range(pl)]
Z1 = ntt.second_NTT(Z1, pl, 0, cnt, True)
Z2 = ntt.second_NTT(Z2, pl, 1, cnt, True)
Z = [ntt.CRT(2, [Z1[i], Z2[i]], ntt.modulo_list) for i in range(pl)]
pnt = 2*m
hap = 0
num = 0
while num < M:
    if Z[pnt] + num < M:
        hap += Z[pnt] * pnt
        num += Z[pnt]
    else:
        hap += pnt * (M - num)
        num = M
    pnt -= 1
print(hap)

Submission status: Runtime Error

Solution code:
class NTT:
    def __init__(self, convolution_rank, binary_level):
        self.cr = convolution_rank
        self.bl = binary_level
        self.modulo_list = [1]*convolution_rank
        self.primal_root_list = [1]*convolution_rank
        self.bin_list = [1]*(binary_level+1)
        self.primal_base_matrix = [[0 for _ in range(binary_level+1)] for __ in range(convolution_rank)]
        self.inverse_base_matrix = [[0 for _ in range(binary_level+1)] for __ in range(convolution_rank)]
        for i in range(binary_level):
            self.bin_list[i+1] = self.bin_list[i] * 2
    
    def doubling(self, n, m, modulo=0):
        y = 1
        tmp = m
        bas = n
        while tmp:
            if tmp % 2:
                y *= bas
                if modulo:
                    y %= modulo
            bas *= bas
            if modulo:
                bas %= modulo
            tmp >>= 1
        return y
    
    def powlimit(self, n):
        y = 1
        cnt = 0
        while y < n:
            y *= 2
            cnt += 1
        return y, cnt
    
    def IsPrime(self, num):
        p = 2
        if num <= 1:
            return False
        while p * p <= num:
            if num % p == 0:
                return False
            p += 1
        return True
    
    def extgcd(self, a, b, c):
        x, y, u, v, k, l = 1, 0, 0, 1, a, b
        while l:
            x, y, u, v = u, v, x - u * (k // l), y - v * (k // l)
            k, l = l, k % l
        if c % k:
            return "No Solution"
        return x * c, y * c
    
    def inved(self, a, modulo):
        x, y = self.extgcd(a, modulo, 1)
        return (x+modulo)%modulo
    
    def root_manual(self):
        for i in range(self.cr):
            r = 1
            flg = True
            while flg:
                fflg = True
                for j in range(self.bl):
                    if self.doubling(r, self.bin_list[j], self.modulo_list[i]) == 1:
                        fflg = False
                if self.doubling(r, self.bin_list[-1], self.modulo_list[i]) != 1:
                    fflg = False
                if fflg:
                    flg = False
                else:
                    r += 1
                    if r >= self.modulo_list[i]:
                        break
            self.primal_root_list[i] = r
    
    def make_prime_root(self):
        cnt = 0
        j = 1
        last = self.bin_list[-1]
        while cnt < self.cr:
            if self.IsPrime(j*last+1):
                flg = True
                r = 1
                while flg:
                    fflg = True
                    for i in range(self.bl):
                        if self.doubling(r, self.bin_list[i], j*last+1) == 1:
                            fflg = False
                    if self.doubling(r, last, j*last+1) != 1:
                        fflg = False
                    if fflg:
                        flg = False
                    else:
                        r += 1
                        if r >= j*last+1:
                            break
                if flg==False:
                    self.modulo_list[cnt] = j*last+1
                    self.primal_root_list[cnt] = r
                    cnt += 1
            j += 2
    
    def make_basis(self):
        for i in range(self.cr):
            for j in range(self.bl):
                tmp = self.doubling(2, self.bl-j)
                self.primal_base_matrix[i][j] = self.doubling(self.primal_root_list[i], tmp, self.modulo_list[i])
                self.inverse_base_matrix[i][j] = self.inved(self.primal_base_matrix[i][j], self.modulo_list[i])
    
    def NTT(self, f, n, idx, depth, inverse=False, surface=True):
        mod = self.modulo_list[idx]
        i2 = self.inved(2, mod)
        if inverse:
            seed = self.primal_base_matrix[idx][depth]
        else:
            seed = self.inverse_base_matrix[idx][depth]
        if n == 1:
            return f
        if n == 2:
            xf = [(f[0] + f[1])%mod, (f[0] - f[1])%mod]
            if inverse * surface:
                xf[0] *= i2
                xf[0] %= mod
                xf[1] *= i2
                xf[1] %= mod
            return xf
        hn = n // 2
        fe = [f[2*i+0] for i in range(hn)]
        fo = [f[2*i+1] for i in range(hn)]
        fe = self.NTT(fe, hn, idx, depth-1, inverse, False)
        fo = self.NTT(fo, hn, idx, depth-1, inverse, False)
        xf = [0 for _ in range(n)]
        grow = 1
        for i in range(hn):
            right = (fo[i] * grow) % mod
            xf[i+0*hn] = (fe[i] + right) % mod
            xf[i+1*hn] = (fe[i] - right) % mod
            grow *= seed
            grow %= mod
        if inverse * surface:
            invn = self.inved(n, mod)
            for i in range(n):
                xf[i] *= invn
                xf[i] %= mod
        return xf
    
    def second_NTT(self, f, n, idx, depth, inverse=False):
        res = [0 for _ in range(n)]
        tmp = [0 for _ in range(n)]
        MOD = self.modulo_list[idx]
        ipl = self.inved(n, MOD)
        for i in range(n):
            bas = 1
            pos = 0
            for j in range(depth, 0, -1):
                pos += bas * ((i>>(j-1)) % 2)
                bas *= 2
            res[i] = f[pos]
        for i in range(depth):
            grow = 1
            seed = inverse * self.primal_base_matrix[idx][i+1] + (1 - inverse) * self.inverse_base_matrix[idx][i+1]
            for k in range(1<<i):
                for j in range(1<<(depth-i-1)):
                    tmp[j*(1<<(i+1))+k+0*(1<<i)] = (res[j*(1<<(i+1))+k] + grow * res[j*(1<<(i+1))+k+(1<<i)]) % MOD
                    tmp[j*(1<<(i+1))+k+1*(1<<i)] = (res[j*(1<<(i+1))+k] - grow * res[j*(1<<(i+1))+k+(1<<i)]) % MOD
                grow *= seed
                grow %= MOD
            for j in range(n):
                res[j] = tmp[j]
        if inverse:
            for i in range(n):
                res[i] *= ipl
                res[i] %= MOD
        return res
    
    def CRT(self, num, a_list, m_list):
        r = a_list[0]
        bas = m_list[0]
        x, y = 0, 0
        for i in range(1, num):
            x, y = self.extgcd(bas, -m_list[i], -a_list[i]-r)
            r += bas * x
            bas *= m_list[i]
        return r % bas

N, M = map(int, input().split())
ntt = NTT(1, 28)
ntt.make_prime_root()
ntt.make_basis()
A = list(map(int, input().split()))
m = max(A)+1
pl, cnt = ntt.powlimit(2*m)
X = [0 for _ in range(pl)]
for i in range(N):
    X[A[i]] += 1
X = ntt.second_NTT(X, pl, 0, cnt)
Z = [X[i]*X[i]%ntt.modulo_list[0] for i in range(pl)]
Z = ntt.second_NTT(Z, pl, 0, cnt, True)
pnt = 2*m
hap = 0
num = 0
while num < M:
    if Z[pnt] + num < M:
        hap += Z[pnt] * pnt
        num += Z[pnt]
    else:
        hap += pnt * (M - num)
        num = M
    pnt -= 1
print(hap)

Submission status: Runtime Error

Solution code:
class NTT:
    def __init__(self, convolution_rank, binary_level):
        self.cr = convolution_rank
        self.bl = binary_level
        self.modulo_list = [1]*convolution_rank
        self.primal_root_list = [1]*convolution_rank
        self.bin_list = [1]*(binary_level+1)
        self.primal_base_matrix = [[0 for _ in range(binary_level+1)] for __ in range(convolution_rank)]
        self.inverse_base_matrix = [[0 for _ in range(binary_level+1)] for __ in range(convolution_rank)]
        for i in range(binary_level):
            self.bin_list[i+1] = self.bin_list[i] * 2
    
    def doubling(self, n, m, modulo=0):
        y = 1
        tmp = m
        bas = n
        while tmp:
            if tmp % 2:
                y *= bas
                if modulo:
                    y %= modulo
            bas *= bas
            if modulo:
                bas %= modulo
            tmp >>= 1
        return y
    
    def powlimit(self, n):
        y = 1
        cnt = 0
        while y < n:
            y *= 2
            cnt += 1
        return y, cnt
    
    def IsPrime(self, num):
        p = 2
        if num <= 1:
            return False
        while p * p <= num:
            if num % p == 0:
                return False
            p += 1
        return True
    
    def extgcd(self, a, b, c):
        x, y, u, v, k, l = 1, 0, 0, 1, a, b
        while l:
            x, y, u, v = u, v, x - u * (k // l), y - v * (k // l)
            k, l = l, k % l
        if c % k:
            return "No Solution"
        return x * c, y * c
    
    def inved(self, a, modulo):
        x, y = self.extgcd(a, modulo, 1)
        return (x+modulo)%modulo
    
    def root_manual(self):
        for i in range(self.cr):
            r = 1
            flg = True
            while flg:
                fflg = True
                for j in range(self.bl):
                    if self.doubling(r, self.bin_list[j], self.modulo_list[i]) == 1:
                        fflg = False
                if self.doubling(r, self.bin_list[-1], self.modulo_list[i]) != 1:
                    fflg = False
                if fflg:
                    flg = False
                else:
                    r += 1
                    if r >= self.modulo_list[i]:
                        break
            self.primal_root_list[i] = r
    
    def make_prime_root(self):
        cnt = 0
        j = 1
        last = self.bin_list[-1]
        while cnt < self.cr:
            if self.IsPrime(j*last+1):
                flg = True
                r = 1
                while flg:
                    fflg = True
                    for i in range(self.bl):
                        if self.doubling(r, self.bin_list[i], j*last+1) == 1:
                            fflg = False
                    if self.doubling(r, last, j*last+1) != 1:
                        fflg = False
                    if fflg:
                        flg = False
                    else:
                        r += 1
                        if r >= j*last+1:
                            break
                if flg==False:
                    self.modulo_list[cnt] = j*last+1
                    self.primal_root_list[cnt] = r
                    cnt += 1
            j += 2
    
    def make_basis(self):
        for i in range(self.cr):
            for j in range(self.bl):
                tmp = self.doubling(2, self.bl-j)
                self.primal_base_matrix[i][j] = self.doubling(self.primal_root_list[i], tmp, self.modulo_list[i])
                self.inverse_base_matrix[i][j] = self.inved(self.primal_base_matrix[i][j], self.modulo_list[i])
    
    def NTT(self, f, n, idx, depth, inverse=False, surface=True):
        mod = self.modulo_list[idx]
        i2 = self.inved(2, mod)
        if inverse:
            seed = self.primal_base_matrix[idx][depth]
        else:
            seed = self.inverse_base_matrix[idx][depth]
        if n == 1:
            return f
        if n == 2:
            xf = [(f[0] + f[1])%mod, (f[0] - f[1])%mod]
            if inverse * surface:
                xf[0] *= i2
                xf[0] %= mod
                xf[1] *= i2
                xf[1] %= mod
            return xf
        hn = n // 2
        fe = [f[2*i+0] for i in range(hn)]
        fo = [f[2*i+1] for i in range(hn)]
        fe = self.NTT(fe, hn, idx, depth-1, inverse, False)
        fo = self.NTT(fo, hn, idx, depth-1, inverse, False)
        xf = [0 for _ in range(n)]
        grow = 1
        for i in range(hn):
            right = (fo[i] * grow) % mod
            xf[i+0*hn] = (fe[i] + right) % mod
            xf[i+1*hn] = (fe[i] - right) % mod
            grow *= seed
            grow %= mod
        if inverse * surface:
            invn = self.inved(n, mod)
            for i in range(n):
                xf[i] *= invn
                xf[i] %= mod
        return xf
    
    def second_NTT(self, f, n, idx, depth, inverse=False):
        res = [0 for _ in range(n)]
        tmp = [0 for _ in range(n)]
        MOD = self.modulo_list[idx]
        ipl = self.inved(n, MOD)
        for i in range(n):
            bas = 1
            pos = 0
            for j in range(depth, 0, -1):
                pos += bas * ((i>>(j-1)) % 2)
                bas *= 2
            res[i] = f[pos]
        for i in range(depth):
            grow = 1
            seed = inverse * self.primal_base_matrix[idx][i+1] + (1 - inverse) * self.inverse_base_matrix[idx][i+1]
            for k in range(1<<i):
                for j in range(1<<(depth-i-1)):
                    tmp[j*(1<<(i+1))+k+0*(1<<i)] = (res[j*(1<<(i+1))+k] + grow * res[j*(1<<(i+1))+k+(1<<i)]) % MOD
                    tmp[j*(1<<(i+1))+k+1*(1<<i)] = (res[j*(1<<(i+1))+k] - grow * res[j*(1<<(i+1))+k+(1<<i)]) % MOD
                grow *= seed
                grow %= MOD
            for j in range(n):
                res[j] = tmp[j]
        if inverse:
            for i in range(n):
                res[i] *= ipl
                res[i] %= MOD
        return res
    
    def CRT(self, num, a_list, m_list):
        r = a_list[0]
        bas = m_list[0]
        x, y = 0, 0
        for i in range(1, num):
            x, y = self.extgcd(bas, -m_list[i], -a_list[i]-r)
            r += bas * x
            bas *= m_list[i]
        return r % bas

N, M = map(int, input().split())
ntt = NTT(1, 27)
ntt.make_prime_root()
ntt.make_basis()
A = list(map(int, input().split()))
m = max(A)+1
pl, cnt = ntt.powlimit(2*m)
X = [0 for _ in range(pl)]
for i in range(N):
    X[A[i]] += 1
X = ntt.second_NTT(X, pl, 0, cnt)
Z = [X[i]*X[i]%ntt.modulo_list[0] for i in range(pl)]
Z = ntt.second_NTT(Z, pl, 0, cnt, True)
pnt = 2*m
hap = 0
num = 0
while num < M:
    if Z[pnt] + num < M:
        hap += Z[pnt] * pnt
        num += Z[pnt]
    else:
        hap += pnt * (M - num)
        num = M
    pnt -= 1
print(hap)

Submission status: Runtime Error

Solution code:
N, M = [int(x) for x in input().split()]
A = sorted([int(x) for x in input().split()], reverse=True) # O(N*logN)
S = [0]*N
S[0] = A[0]
for i in range(1, N):
    S[i] = S[i - 1] + A[i]

#Aの要素のペアであって和がxより大になるものの個数とそれらの和
def f(x, A): # O(N*logN)
    for i in range(1, N):
        S[i] = S[i - 1] + A[i]
    num = 0
    tot = 0
    for i in range(N):
        low = 0
        high = N - 1
        while low < high:
            mid = (low + high + 1) // 2
            if A[i] + A[mid] > x:
                low = mid
            else:
                high = mid - 1
        if A[i] + A[low] > x:
            num += low + 1
            tot += A[i] * (low + 1) + S[low]
    
    return (num, tot)

low = 0
high = A[0] + A[1]
while low < high:
    mid = (low + high) // 2
    if f(mid, A)[0] < M:
        high = mid
    else:
        low = mid + 1

K = low #最適解における最小の幸福度上昇値
M0 = f(K, A)[0] #幸福度がKより大となるペアの数
S0 = f(K, A)[1] #幸福度がKより大となるペアをすべて足した和
ans = S0 + (M - M0) * K
print(ans)

Submission status: Runtime Error

Solution code:
import sys

# 価値 x 以上の握手を全て行うとする
# 握手回数が決まる。左手ごとに集計できる。

import numpy as np

N, M = map(int, sys.stdin.readline().rstrip().split())
A = np.array(int(x) for x in sys.stdin.readline().rstrip().split())

A.sort()  # A を昇順に並び替える


# x より価値が高くなる握手の総数
#  -> a に対して、相手は x - a が必要
#   -> A = [a0, a1, a2, ...] に対して、[x-a0の開始位置, x-a1の開始位置, x-a2の開始位置,...] を探す
def shake_cnt(x):
    # x 以上の握手を全て行うとして、握手の回数を数える行わない握手を数える
    X = np.searchsorted(A, x - A)  # ソートされたリストA に対して、x - A が挿入される位置を返す
    return N * N - X.sum()  # index の和は「小さい方」を指しているので、引く


# 二分探索  握手の回数が
left = 0  # 握手の回数がM以上
right = 10 ** 6  # 握手の回数がM未満 <- まだ足りない側
while 1 < right - left:  # left と right が隣り合ったら終了
    x = (left + right) // 2
    if shake_cnt(x) >= M:
        left = x
    else:
        right = x

# M回には届かないけれど、M回に限りなく近くなる価値 x = right
X = np.searchsorted(A, right - A)  # 各a in A が握手を行わない人数のArray
shake = N * N - X.sum()  # ここまでの握手回数

Acum = np.zeros(N + 1, np.int64)  # 人数 -> 累積和
# Acum = [a0, a0+a1, a0+a1+a2, ....] を計算して、ai さんが k番目以降と握手をすれば
Acum[1:] = np.cumsum(A)  # cumsum [1,2,3,4,5,6,...] -> [1,3,6,10,15,21,...]

# 右手の幸福度は ai * ((a0+a1+...a(n-1)) - (a(k-1)+...+a(n-1))) を得られる
# 左手の幸福度は 回数だけで決まる（何人と握手とするか）
happy = (Acum[-1] - Acum[X]).sum() + (A * (N - X)).sum()

happy += (M - shake) * left  # 不足分。1回の価値は M とわかっている

print(happy)

Submission status: Runtime Error

Solution code:
n, m = map(int, input().split())
a = sorted(list(map(int, input().split())), reverse=True)
s = a[0]
t = 0
j = 1
ret = 0
for i in range(n):
    s -= a[i]
    ret += a[i] * 2
    m -= 1
    while j < n and (i >= n - 1 or a[i + 1] + a[i + 1] <= a[i] + a[j]):
        if m < i * 2:
            for k in range(i):
                if m > 0:
                    ret += a[j] + a[k]
                    m -= 1
                if m > 0:
                    ret += a[j] + a[k]
                    m -= 1
        else:
            ret += a[j] * i * 2 + t * 2
            m -= i * 2
        s += a[j]
        j += 1
    num = (j - i - 1) * 2
    if m <= num:
        for k in range(i + 1, j):
            if m > 0:
                ret += a[i] + a[k]
                m -= 1
            if m > 0:
                ret += a[i] + a[k]
                m -= 1
        break
    m -= num
    ret += s * 2 + a[i] * num
    t += a[i]
    print(a, m, num, s, i, j, ret)
print(ret)

Submission status: Runtime Error

Solution code:
from cmath import pi, exp

def convolve(a, b):
    """配列a, bの畳み込みをFFT(fast Fourier transform)によって計算する
    res[k] = sum{i=0...n}a[i]*b[k-i]
    """
    def fft(a, inverse=False):
        # バタフライ演算用の配置入れ替え
        for i in range(n):
            j = 0
            for k in range(h):
                j |= (i >> k & 1) << (h - 1 - k)
            if i < j:
                a[i], a[j] = a[j], a[i]
 
        # バタフライ演算
        sign = 1.0 if inverse else -1.0
        b = 1
        while b < n:
            for j in range(b):
                w = exp(sign * 2.0j * pi / (2.0 * b) * j)
                for k in range(0, n, b*2):
                    s = a[j + k]
                    t = a[j + k + b] * w
                    a[j + k] = s + t
                    a[j + k + b] = s - t
            b *= 2
 
        # 逆変換時には配列のサイズで割る
        if inverse:
            for i in range(n):
                a[i] /= n
        return a
 
    def ifft(a):
        return fft(a, inverse=True)
 
    n = 1 << (len(a) + len(b) - 1).bit_length()
    h = n.bit_length() - 1
    a += [0] * (n - len(a))
    b += [0] * (n - len(b))
    
    # in-placeで計算を進めることに注意
    fft(a), fft(b)
    for i in range(n):
        a[i] *= b[i]
    ifft(a)
 
    res = [int(abs(v) + 0.5) for i, v in enumerate(a)]
    return res[1:]

import sys
input = sys.stdin.readline
n,m = map(int, input().split())
a = list(map(int, input().split()))
 
cnt = [0] * (10010)
for i in a:
    cnt[i] += 1 

c = convolve(cnt[0:], cnt[0:])
ans = 0
for i in range(len(c))[::-1]:
    if c[i] > 0:
      p = min(m,c[i])
      m -= p
      ans += i*p+p
      if m == 0:
        break
print(ans)

Submission status: Runtime Error

Solution code:
from cmath import pi, exp

def convolve(a, b):
    """配列a, bの畳み込みをFFT(fast Fourier transform)によって計算する
    res[k] = sum{i=0...n}a[i]*b[k-i]
    """
    def fft(a, inverse=False):
        # バタフライ演算用の配置入れ替え
        for i in range(n):
            j = 0
            for k in range(h):
                j |= (i >> k & 1) << (h - 1 - k)
            if i < j:
                a[i], a[j] = a[j], a[i]
 
        # バタフライ演算
        sign = 1.0 if inverse else -1.0
        b = 1
        while b < n:
            for j in range(b):
                w = exp(sign * 2.0j * pi / (2.0 * b) * j)
                for k in range(0, n, b*2):
                    s = a[j + k]
                    t = a[j + k + b] * w
                    a[j + k] = s + t
                    a[j + k + b] = s - t
            b *= 2
 
        # 逆変換時には配列のサイズで割る
        if inverse:
            for i in range(n):
                a[i] /= n
        return a
 
    def ifft(a):
        return fft(a, inverse=True)
 
    n = 1 << (len(a) + len(b) - 1).bit_length()
    h = n.bit_length() - 1
    a += [0] * (n - len(a))
    b += [0] * (n - len(b))
    
    # in-placeで計算を進めることに注意
    fft(a), fft(b)
    for i in range(n):
        a[i] *= b[i]
    ifft(a)
 
    res = [int(abs(v) + 0.5) for i, v in enumerate(a)]
    return res[1:]

import sys
input = sys.stdin.readline
n,m = map(int, input().split())
a = list(map(int, input().split()))
 
cnt = [0] * (10001)
for i in a:
    cnt[i] += 1 

c = convolve(cnt[0:], cnt[0:])
ans = 0
for i in range(len(c))[::-1]:
    if c[i] > 0:
      p = min(m,c[i])
      m -= p
      ans += i*p+p
      if m == 0:
        break
print(ans)

Submission status: Runtime Error

Solution code:
from cmath import pi, exp

def convolve(a, b):
    """配列a, bの畳み込みをFFT(fast Fourier transform)によって計算する
    res[k] = sum{i=0...n}a[i]*b[k-i]
    """
    def fft(a, inverse=False):
        # バタフライ演算用の配置入れ替え
        for i in range(n):
            j = 0
            for k in range(h):
                j |= (i >> k & 1) << (h - 1 - k)
            if i < j:
                a[i], a[j] = a[j], a[i]
 
        # バタフライ演算
        sign = 1.0 if inverse else -1.0
        b = 1
        while b < n:
            for j in range(b):
                w = exp(sign * 2.0j * pi / (2.0 * b) * j)
                for k in range(0, n, b*2):
                    s = a[j + k]
                    t = a[j + k + b] * w
                    a[j + k] = s + t
                    a[j + k + b] = s - t
            b *= 2
 
        # 逆変換時には配列のサイズで割る
        if inverse:
            for i in range(n):
                a[i] /= n
        return a
 
    def ifft(a):
        return fft(a, inverse=True)
 
    n = 1 << (len(a) + len(b) - 1).bit_length()
    h = n.bit_length() - 1
    a += [0] * (n - len(a))
    b += [0] * (n - len(b))
    
    # in-placeで計算を進めることに注意
    fft(a), fft(b)
    for i in range(n):
        a[i] *= b[i]
    ifft(a)
 
    res = [int(abs(v) + 0.5) for i, v in enumerate(a)]
    return res[1:]

import sys
input = sys.stdin.readline
n,m = map(int, input().split())
a = list(map(int, input().split()))
 
cnt = [0] * (10001)
for i in a:
    cnt[i] += 1 

c = convolve(cnt[0:], cnt[0:])
ans = 0
for i in range(20001)[::-1]:
    if c[i] > 0:
      p = min(m,c[i])
      m -= p
      ans += i*p+p
      if m == 0:
        break
print(ans)