[AtCoder] F - K-th Largest Triplet

F - K-th Largest Triplet

  • Time :
  • Space :
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
import sys
import heapq
from collections import Counter

def freq(arr):
counter = Counter(arr)
items = sorted(counter.items(), reverse=True)
return items

def solve(A, B, C, n, k):
fa = freq(A)
fb = freq(B)
fc = freq(C)

pq = []
vis = set()

def push(i, j, l):
if (i, j, l) in vis:
return
if i == len(fa) or j == len(fb) or l == len(fc):
return
vis.add((i, j, l))
a = fa[i][0]
b = fb[j][0]
c = fc[l][0]
val = a*b + b*c + c*a
heapq.heappush(pq, (-val, i, j, l))

push(0, 0, 0)

while True:
neg_val, i, j, l = heapq.heappop(pq)
val = -neg_val
cnt = fa[i][1] * fb[j][1] * fc[l][1]
k -= cnt
if k <= 0:
return val
push(i+1, j, l)
push(i, j+1, l)
push(i, j, l+1)

def main():
input = sys.stdin.readline
n, k = map(int, input().split())
A = list(map(int, input().split()))
B = list(map(int, input().split()))
C = list(map(int, input().split()))
print(solve(A, B, C, n, k))

if __name__ == "__main__":
main()
Author: Song Hayoung
Link: https://songhayoung.github.io/2025/11/23/PS/AtCoder/abc391-f/
Copyright Notice: All articles in this blog are licensed under CC BY-NC-SA 4.0 unless stating additionally.