ARC 074 D 3N Numbers【Python】

問題はこちら。

import sys
readline = sys.stdin.readline

def main():
   from heapq import heapify, heappop, heappush
   
   N = int(readline())
   a = list(map(int, readline().split()))
   
   red = a[:N]
   blue = list(map(lambda x: -x, a[2*N:]))
   ans_red = [0]*(N + 1); ans_blue = [0]*(N + 1)
   heapify(red); heapify(blue)
   ans_red[0], ans_blue[0] = sum(red), sum(blue)
   for k in range(1, N + 1):
       heappush(red, a[k + N - 1])
       ans_red[k] = ans_red[k - 1] + a[k + N - 1]
       sub_red = heappop(red)
       ans_red[k] -= sub_red
       heappush(blue, -a[2*N - k])
       ans_blue[k] = ans_blue[k - 1] - a[2*N - k]
       sub_blue = heappop(blue)
       ans_blue[k] -= sub_blue
   
   ans = -float('inf')
   for k in range(N + 1):
       ans = max(ans, ans_red[k] + ans_blue[N - k])
   print(ans)
   
if __name__ == "__main__":
   main()