F - K-th Largest Triplet Editorial by en_translator
Sort each of \(A\), \(B\), and \(C\) in descending order. Also, let \(f(i,j,k)=A_iB_j+B_jC_k+C_kA_i\).
Let us try enumerating the top \(K\), as \(K\) is as small as \(K\leq 2\times 10^5\). Importantly,
\[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).\]
This implies that if we enumerate the values in descending order, \(f(i,j,k)\) always comes before \(f(i+1,j,k),f(i,j+1,k)\), and \(f(i,j,k+1)\).
Thus, the problem can be solved by the following algorithm:
- Prepare a queue \(Q\). Insert \((f(1,1,1),1,1,1)\) to \(Q\).
- Repeat the following \(K\) times.
- Let \((val,i,j,k)\) be the maximum element in \(Q\). Remove \((val,i,j,k)\), and insert \((f(i+1,j,k),i+1,j,k),(f(i,j+1,k),i,j+1,k)\), and \((f(i,j,k+1),i,j,k+1)\) to \(Q\) (if absent).
The time complexity is \(O(N\log N+K\log K)\).
(Since Python’s priority queue manages minimum vales, we multiply \(f(i,j,k)\) by \(-1\) before inserting it to \(\mathrm{hq}\), and multiply it back by \(-1\) when printing in Line 22.)
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: