Advertisement
Not a member of Pastebin yet?
Sign Up,
it unlocks many cool features!
- """
- Learning task
- https://contest.yandex.ru/contest/40146/problems/E/
- Calculate the total utility of the `k` most useful potions (consisting of one or two ingredients)
- I/O examples (first - n, k; the second - is the usefulness of `n` potions):
- 5 5
- -2 3 -5 5 1
- =
- 26
- 2 1
- -1 1
- =
- 1
- """
- import sys
- import bisect
- def main():
- n, k = map(int, input().split())
- line = sys.stdin.readline().rstrip()
- a = list(map(int, line.split()))
- a.sort()
- a_sum = [0] + a.copy()
- for i in range(1, len(a_sum)):
- a_sum[i] = a_sum[i] + a_sum[i - 1]
- if n == 1:
- print(a[-1])
- return
- def f(x, calc=False):
- cur_idx = bisect.bisect_right(a, x)
- res = n - min(cur_idx, n)
- res_sum = 0
- if calc:
- res_sum = a_sum[-1] - a_sum[min(cur_idx, n)]
- cur_r = len(a) - 1
- for idx, v in enumerate(a):
- cur_x = x - v
- if a[-1] <= cur_x:
- continue
- while cur_r > -1 and a[cur_r] > cur_x and cur_r > idx:
- cur_r -= 1
- new_idx = max(cur_r, idx)
- cur_res = n - new_idx - 1
- res += cur_res
- if calc:
- sum_part = a_sum[-1] - a_sum[min(new_idx+1, n)]
- temp_sum = sum_part + v * (n - new_idx - 1)
- res_sum += temp_sum
- return res, res_sum
- l = min(a[0], a[0] + a[1]) - 1
- r = max(2 * a[-1], 2 * sum(a[-2:])) + 1
- while l < r:
- x = (l + r) // 2
- res_f, _ = f(x)
- if res_f > k:
- l = x + 1
- else:
- r = x
- cur_k, cur_sum = f(l, True)
- print(cur_sum + l * (k - cur_k))
- if __name__ == '__main__':
- main()
Advertisement
Add Comment
Please, Sign In to add comment
Advertisement