Official

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:

  1. Prepare a queue \(Q\). Insert \((f(1,1,1),1,1,1)\) to \(Q\).
  2. 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: