3N Numbers

priority_queueの問題を解いてみた。

atcoder.jp

f:id:poohandduffy:20200915041532p:plain

この問題を言い換えると、真ん中のN個の数字に仕切りを置いて、その左側を最大化、右側を最小化するようにN個ずつ選んで、左の和から右の和を引けばよいことになる。
ただし、仕切りのパターンがN+1個あり、そのたびに最大化/最小化の作業を繰り返すと計算量がN^2になってしまう。

そこで、まずは、K番目に仕切りを置いたときに、左側を最大化することだけを考えることにする。
仕切りが一番左端にあるとき、左側の最大値は、sum(左からN個)となる。そこから右に1個仕切りを動かすと、どうなるか。
左側N個に、中央の左端を加えて、その最小値を間引けばよい。この操作を繰り返すと、順々に仕切りを動かした場合の最大値が求まる。

heapq.heappush(l, m[k]) #真ん中のk番目の値をヒープに追加する。
left[1+k] = left[k] + m[k] - heapq.heappop(l) #新たな合計値は、動かす前の合計値から、加えた真ん中の値をたし、さっ引いたヒープの最小値を引けばよい。

ヒープの挿入/取り出しはlogNなので、N*2logN
左右両方で行うため、4NlogNということか。

仕切りを右に動かして、幅を増やしていくのは簡単だが、逆に減らして行くことはできないのがキーポイントです。
なので、右側でも同様に行えばいいが、独立に行う事になる。
f:id:poohandduffy:20200915043802p:plain
つまり、左側で左からK番目の仕切りの場合を計算しているときに、右側では右からK番目の仕切りの場合を計算していることになるため、それらの最大値、最小値を一度記録しなければならず、ループの中で差を計算することができない。

import heapq

N = int(input())
an = list(map(int, input().split()))

left = [0] * (N+1) #仕切りを0からNまで動かしたときの、仕切りの左側からN個選ぶ時の最大値
right = [0] * (N+1)

left[0] = sum(an[:N])
right[-1] = -sum(an[2*N:])

l = an[:N]
heapq.heapify(l)
m = an[N:2*N]
r = [-i for i in an[2*N:]]
heapq.heapify(r)

for k in range(N):
    heapq.heappush(l, m[k])
    left[1+k] = left[k] + m[k] - heapq.heappop(l)

    heapq.heappush(r, -m[-1-k])
    right[-2-k] = right[-1-k] - heapq.heappop(r) - m[-1-k]

ans = -float('inf')
for i in range(N+1):
    ans = max(ans, left[i]+right[i])

print(ans)