Official

F - K-th Largest Triplet Editorial by toam

解説1

\(A, B, C\) をそれぞれ降順ソートしておきます. \(f(i,j,k)=A_iB_j+B_jC_k+C_kA_i\) とします.

\(K\)\(K\leq 5\times 10^5\) と小さいことを利用して,上位 \(K\) 個を列挙することを考えます.重要な事実として

\[f(i,j,k)\geq f(i+1,j,k),f(i,j,k)\geq f(i,j+1,k),f(i,j,k)\geq f(i,j,k+1)\]

が成り立ちます.よって,大きい方から値を列挙していったときに,必ず \(f(i,j,k)\) が列挙された後に \(f(i+1,j,k),f(i,j+1,k),f(i,j,k+1)\) が列挙されることが分かります.

これを踏まえれば,以下のアルゴリズムで解くことができます.

  1. 優先度付きキュー \(Q\) を用意し, \(Q\)\((f(1,1,1),1,1,1)\) を追加する.
  2. 以下を \(K\) 回繰り返す.
    • \(Q\) の最大値を \((val,i,j,k)\) とする.\(Q\) から \((val,i,j,k)\) を削除したうえで,\(Q\)\((f(i+1,j,k),i+1,j,k),(f(i,j+1,k),i,j+1,k),(f(i,j,k+1),i,j,k+1)\) を(まだ \(Q\) に追加していなかったら)追加する.

計算量は \(O(N\log N+K\log K)\) です.

(Python の優先度付きキューは最小値を管理するため,以下の実装例では \(\mathrm{hq}\) には \(f(i,j,k)\)\(-1\) 倍して追加し,22 行目の出力時に \(-1\) 倍することで元に戻しています.)

from heapq import heappop, heappush

N, K = map(int, input().split())
A = sorted(list(map(int, input().split())), reverse=True)
B = sorted(list(map(int, input().split())), reverse=True)
C = sorted(list(map(int, input().split())), reverse=True)

s = set()
hq = []


def add(i, j, k):
    if i < N and j < N and k < N and (i, j, k) not in s:
        s.add((i, j, k))
        heappush(hq, (-(A[i] * B[j] + B[j] * C[k] + C[k] * A[i]), i, j, k))


add(0, 0, 0)
for t in range(K):
    val, i, j, k = heappop(hq)
    if t == K - 1:
        print(-val)
    add(i + 1, j, k)
    add(i, j + 1, k)
    add(i, j, k + 1)

posted:
last update: