aacord’s memo

abcを中心にpythonで解いた問題のメモ、整理をしています。緑になった。

abc 149 E - Handshake (python)

M回握手をして出来る幸福度の最大値を、1回の握手で幸福度がx以上上がるような握手の回数がM回出来るxの最大値を求める問題と考えれば、二分探索でxを求めて、左手 i を固定したときにx以上の握手ができる手の数は、
a を昇順ソートしておけば k = n - bisect_left(a,x-i) で求めることができ、その合計は i*k + sum(a[n-k:]) となるので、事前に累積和のリストを作っておけば O(1)で左手を固定したときの答えが求まる。
ただし、サンプル2のように X 以上の上がる握手の回数 k の総和がMを超えることがあるので、
k の総和が M を超えた分だけxを引く必要がある。(二分探索の結果、M回握手できる下界の最大値はxなので、M を超えて数えている幸福度は必ず x になる。もし、x+1も含まれているのなら、少なくともx+1が二分探索で条件をみたす。)

import sys
read = sys.stdin.buffer.read
readline = sys.stdin.buffer.readline
readlines = sys.stdin.buffer.readlines
from bisect import bisect_left

n,m = map(int,readline().split())
a = list(map(int,readline().split()))

b = [0]*(n+1)
a.sort()
for i in range(n):
  b[i+1] = b[i]+a[i]

def is_ok(arg):
    c = 0
    f = False
    for i in a:
      if c >= m:
        break
      c += (n-bisect_left(a,arg-i))
    if c >= m:
      f = True
    return f

def bisect_ok(ng, ok):
    while (abs(ok - ng) > 1):
        mid = (ok + ng) // 2
        if is_ok(mid):
            ok = mid
        else:
            ng = mid
    return ok

x = bisect_ok(2*a[-1]+1,1)
ans = 0
cnt = 0
for i in a:
  ind = n-bisect_left(a,x-i)
  cnt += ind
  ans += i*ind+(b[-1]-b[n-ind])
if cnt > m:
  ans -= x*(cnt-m)
print(ans)