[백준/Python]2295: 세 수의 합

문제

https://www.acmicpc.net/problem/2295

백준 문제집 - 0x13강 - 이분탐색

 

알고리즘: 이분 탐색, 해시를 사용한 집합과 맵

---

풀이

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

우리가 x번째 수, y번째 수, z번째 수를 더해서 k번째 수를 만들었다라고 하자. 위의 예제에서 2+3+5=10의 경우는 x, y, z, k가 차례로 1, 2, 3, 4가 되며, 최적해의 경우는 2, 3, 4, 5가 된다. k번째 수가 최대가 되도록 하는 것이 목적이다. x, y, z, k가 서로 같아도 된다. 이때, k번째 수를 출력하면 된다.

 

집합 내의 원소를 세 번 골라서 더한값이 집합에 존재하는경우를 찾아야한다.

설명에 "x, y, z, k가 서로 같아도 된다" 라고 나와있는데, 갯수는 신경쓰지 않아도 됨을 명심하자.

예를들어 집합이 {1, 1, 2, 6} 로 구성된경우 최적해는 2 + 2 + 2 = 6 이 되는것이다.

 

사실 시간초과의 쓴맛을 몇번 봤던 문제다.

처음 시도했던 방법은 집합의 제일 큰값을 k로 설정하고 하나씩 줄여나가면서, 이중for문으로 x, y를 선택 후 이분탐색을 통해 z를 찾는것이었다.

여기서 큰 틀은 바꾸지 않은 채 디테일만 손봤었고 결국 infinity TLE...♾️

 

이렇게 몇 번 고생하다가 찾은 해답은 이러하다.

1. 두 수의 합을 set()에 미리 저장해둔다.
2. 집합 U의 가장 큰 값부터 내림차순으로 순회한다. → k
3. 각 k마다 집합의 원소를 다시 순회한다. → z
4. k - z 의 값을 이분탐색을 통해 찾는다. 이때 정렬된 set()을 사용한다.

 

k는 가장 큰 값부터 시작하므로 제일 먼저 발견되는 k - z의값이 곧 답이 되는것이다.

또한 주어지는 원소들은 모두 자연수이므로 k - z 가 음수라면 이분탐색을 실행하지 않고 바로 넘어가면된다.

 

보면 첫번째 방법과 크게 다르지 않다.

단, 첫번째는 x, y를 매번 더해준 뒤 z를 찾는것이었고 두번째는 z를 기준으로 x+y를 찾는것이므로 시간 차이가 발생하게 된다.

역시 사람은 효율적으로 살아야한다.

 

전체 코드

# 메모리: 71040KB / 시간: 760ms
from sys import stdin


input = stdin.readline

N = int(input())
arr = sorted(int(input()) for _ in range(N))

two_sum = set()  # 두 수를 더한 값들을 미리 저장해둠.
for i in range(N):
    for j in range(i, N):
        two_sum.add(arr[i] + arr[j])

two_sum = sorted(two_sum)


def binary_search(target):
    start, end = 0, len(two_sum)-1

    while start <= end:
        mid = (start + end) // 2

        if two_sum[mid] == target:
            return True
        elif two_sum[mid] < target:
            start = mid + 1
        else:
            end = mid - 1
    return False


def finding():
    for k in range(N-1, -1, -1):
        d = arr[k]
        for i in range(N):
            c = d - arr[i]  # c = 두 수를 더한 값
            if c <= 0:
                continue
            # 만약 c가 two_sum 리스트에 존재한다면 조건을 만족하므로 바로 반환
            if binary_search(c):
                return d


print(finding())