Official

I - K-th Largest Triplet Editorial by toam

解説2

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

答えを二分探索します.判定問題として \(f(i,j,k)\geq mid\) を満たす \((i,j,k)\)\(K\) 個あるか?が高速に求められれば良いです.

判定問題を愚直で解こうとすると \(O(N^3)\) かかり間に合いませんが,枝刈りをすることで計算量を落とすことができます.

\[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)\]

が成り立つことを利用すると,次のような判定関数のコードを書くことができます.

bool check(long long mid) {
    int cnt = 0;
    for (int i = 1; i <= N; i++) {
        if (f(i, 1, 1) < mid) break;  // もし f(i, 1, 1) が mid より小さいならば i のループを打ち切る
        for (int j = 1; j <= N; j++) {
            if (f(i, j, 1) < mid) break;  // もし f(i, j, 1) が mid より小さいならば j のループを打ち切る
            for (int k = 1; k <= N; k++) {
                if (f(i, j, k) < mid) break;  // もし f(i, j, k) が mid より小さいならば k のループを打ち切る
                cnt += 1;
                if (cnt == K) return true;  // もし cnt が K ならば true を返す
            }
        }
    }
    return false  // false を返す
}

これにより,判定関数の中でループが回る回数を \(O(K)\) 回に減らすことができます.よって,この問題を計算量 \(O(N\log N+K(\log \max A+\log \max B+\log \max C))\) で解くことができます.

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)


def f(i, j, k):
    return A[i] * B[j] + B[j] * C[k] + C[k] * A[i]


def calc(mid):
    cnt = 0
    for i in range(N):
        if f(i, 0, 0) < mid:
            break
        for j in range(N):
            if f(i, j, 0) < mid:
                break
            for k in range(N):
                if f(i, j, k) >= mid:
                    cnt += 1
                    if cnt == K:
                        return True
                else:
                    break
    return False


ok, ng = 0, 3 * 10**18 + 1
while ng - ok > 1:
    mid = (ok + ng) // 2
    if calc(mid):
        ok = mid
    else:
        ng = mid

print(ok)

posted:
last update: