두 개의 정수 리스트 p와 q가 주어졌을 때, 두 리스트의 모든 요소를 서로 곱한 결과들 중에서 k번째로 큰 값을 찾아야 하는 문제입니다.
예를 들어, 입력이 p = [2, 5], q = [6, 8], k = 2라면 출력은 16이 됩니다.
곱셈 결과를 살펴보면 다음과 같습니다: 2 × 6 = 12, 2 × 8 = 16, 5 × 6 = 30, 5 × 8 = 40. 이 결과들을 내림차순으로 정렬하면 40, 30, 16, 12이며, 0부터 시작하는 인덱스 기준으로 2번째로 큰 값은 16입니다.
해결 접근 방식
모든 곱을 계산한 뒤 정렬하는 방법도 있지만, 최소 힙(min-heap)을 활용하면 더 효율적으로 상위 k개의 값만 유지할 수 있습니다. 해결 단계는 다음과 같습니다.
- 리스트 p를 오름차순으로 정렬합니다.
- 리스트 q를 오름차순으로 정렬합니다.
- k에 1을 더해 인덱스 보정을 합니다 (k := k + 1).
- 빈 리스트 형태의 힙(heap)을 생성합니다.
- q의 각 요소(elem)에 대해 다음을 반복합니다:
- elem이 0 이상인 경우: p를 마지막 인덱스부터 0까지 역순으로 순회하며 cd = elem × p[i]를 계산합니다. 힙이 비어 있지 않고 힙의 크기가 k와 같으며 cd ≤ heap[0]이라면 반복문을 종료하고, 그렇지 않으면 cd를 힙에 삽입합니다. 힙의 길이가 k보다 커지면 가장 작은 값을 제거합니다.
- elem이 음수인 경우: p를 처음부터 끝까지 순서대로 순회하며 동일하게 cd = elem × p[i]를 계산합니다. 조건이 충족되면 반복문을 종료하고, 아니면 cd를 힙에 삽입한 뒤 힙의 길이가 k를 초과하면 가장 작은 값을 제거합니다.
- 최종적으로 heap[0]을 반환합니다.
음수 처리 로직의 핵심
곱셈에서 부호의 성질 때문에 탐색 순서가 달라집니다. 양수와 곱할 때는 p의 큰 값부터 순회해야 큰 곱을 먼저 얻고, 음수와 곱할 때는 p의 작은 값(절대값이 큰 음수)부터 순회해야 큰 곱을 먼저 얻습니다. 이 덕분에 불필요한 계산을 조기에 중단(break)할 수 있습니다.
예제 구현
다음 구현을 통해 더 잘 이해해 보겠습니다.
from heapq import heappush, heappop
def solve(p, q, k):
p = sorted(p)
q = sorted(q)
k += 1
heap = []
for elem in q:
if elem >= 0:
for i in range((len(p) - 1), -1, -1):
cd = elem * p[i]
if heap and len(heap) == k and cd <= heap[0]:
break
heappush(heap, cd)
if len(heap) > k:
heappop(heap)
else:
for i in range(len(p)):
cd = elem * p[i]
if heap and len(heap) == k and cd <= heap[0]:
break
heappush(heap, cd)
if len(heap) > k:
heappop(heap)
return heap[0]
print(solve([2, 5], [6, 8], 2))입력
[2, 5], [6, 8], 2
출력
16
이 알고리즘은 힙의 크기를 항상 k 이하로 유지하기 때문에, 모든 곱을 저장하고 정렬하는 O(n·m·log(n·m)) 방식보다 메모리 측면에서 유리하며, 조기 종료 조건을 통해 실제 연산량도 크게 줄일 수 있습니다.