Official
I - K-th Largest Triplet Editorial
by
解説1
I - 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)\) が列挙されることが分かります.
これを踏まえれば,以下のアルゴリズムで解くことができます.
- 優先度付きキュー \(Q\) を用意し, \(Q\) に \((f(1,1,1),1,1,1)\) を追加する.
- 以下を \(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: