[2295] 세 수의 합 - Python

/* 문제 */
N(5 ≤ N ≤ 1,000)개의 자연수들로 이루어진 집합 U가 있다. 이 중에서 적당히 세 수를 골랐을 때, 그 세 수의 합 d도 U안에 포함되는 경우가 있을 수 있다. 이러한 경우들 중에서, 가장 큰 d를 찾으라.
예를 들어 {2, 3, 5, 10, 18}와 같은 집합이 있다고 하자. 2+3+5 = 10이 되고, 이 수는 집합에 포함된다. 하지만 3+5+10 = 18이 되고, 이 경우가 세 수의 합이 가장 커지는 경우이다.

풀이 방법 자체는 간단하게 떠올릴 수 있지만 시간 제한이 있기 때문에 효율적으로 푸는 방법을 고민해야 한다.

 

처음에는 itertools를 이용해 리스트에서 세 수를 뽑은 후 더해 그 합이 리스트에 포함되는지를 확인했다. 불필요한 계산을 줄이기 위해 세 수를 뽑을 때 큰 숫자부터 뽑고, 첫 답을 찾았을 때 코드를 바로 종료했다.

# 시간 초과 코드
import sys
from itertools import combinations_with_replacement
input = sys.stdin.readline

n = int(input())
num = [int(input()) for _ in range(n)]
for i in combinations_with_replacement(num[::-1], 3):
    tem = sum(i)
    if tem in num:
        print(tem)
        exit()

combination을 사용하면 리스트의 요소들로 튜플들을 만들어 준다. 이해를 잘못해서 반복문을 돌 때마다 튜플을 하나씩 머만들어 준다고 생각했는데, 미리 튜플 리스트를 만들고 그 안의 요소를 뽑아오는 것 같다. combination의 시간 복잡도는 O(n!)으로 높은 편이다. 따라서 반복문 중간에 exit한다고 해도 시간 초과가 뜰 수 밖에 없다.

 

# 정답 코드
import sys
input = sys.stdin.readline

n = int(input())
num = [int(input()) for _ in range(n)]
num_sum = set()
for x in num:
    for y in num:
        num_sum.add(x + y)
num.sort()
for k in range(n-1, -1, -1):
    for y in range(k+1):
        if num[k] - num[y] in num_sum:
            print(num[k])
            exit()

다음으로 생각한 방법은 x+y+z = k이므로 x+y = k-z임을 이용하는 것이다.

x+y를 계산하여 집합에 저장한다. 집합에 저장하는 과정에서 중복 값도 제거할 수 있다.

숫자 리스트를 정렬한 후 뒤에서부터 k값을 찾고, 가장 큰 k값을 찾는 즉시 종료하여 불필요한 연산을 줄였다.