[Python/코테] 백준 2295 세 수의 합

2295 세 수의 합

문제 및 조건 설명: https://www.acmicpc.net/problem/2295

알고리즘 설계

 a[i] + a[j] + a[k] = a[l]을 만족하는 a[l] 중에서 최댓값을 구하기

 

시간 복잡도를 줄여나가는 풀이를 떠올리는 과정

1. $O(N^4)$ - i, j, k, l에 대한 4중 for문
2. $O(N^3lgN)$ - i, j, k에 대한 3중 for문을 돌리고, 배열 a 안에 a[i]+a[j]+a[k]가 있는지 이분탐색 ( N = 1000이라 시간 초과)
3. $O(N^2lgN)$ - 미리 a에서 두 원소의 합을 다 묶어놓은 배열 two를 만들고, two[m] + a[k] = a[l] 로 식을 변경.  k, l에 대한 2중 for문을 돌리고 a[l]-a[k]가 배열 two 안에 있는지 확인하는 방법 
   ( two의 길이는 $N2$인데 $lg(N^2)$은 $2lgN$이어서 $O(N^2lg(N^2)) = O(N^2 * 2lgN) = O(N^2lgN)$)

a[l]-a[k]가 배열 two 안에 있는지 확인할 때, 파이썬의 in 연산자를 사용하면 시간 복잡도가 $O(N)$이 된다. 즉 시간 초과 발생.

이진 탐색은 $O(logN)$이므로, 효율적인 코드를 위해 이진 탐색을 구현하도록 한다.

알고리즘 구현

1. 이진 탐색 직접 구현

import sys

n = int(input())
list = [int(sys.stdin.readline()) for _ in range(n)]
list.sort()

sum = []

for i in range(n):
    for j in range(i, n):
        sum.append(list[i] + list[j])
sum.sort()


def check(i):
    st = 0
    en = len(sum) - 1

    while st <= en:
        mid = (st+en)//2
        if sum[mid] == i:
            return True
        elif sum[mid] > i:
            en = mid - 1
        else:
            st = mid + 1

    return False


def search():
    for i in reversed(range(n)):
        for j in range(n):
            if check(list[i] - list[j]):
                print(list[i])
                return


search()

2. 라이브러리 사용

import sys
from bisect import bisect_left, bisect_right

n = int(input())
list = [int(sys.stdin.readline()) for _ in range(n)]
list.sort()

sum = []

for i in range(n):
    for j in range(i, n):
        sum.append(list[i] + list[j])
sum.sort()


def func():
    for i in range(n-1, 0, -1):
        for j in range(i):
            diff = list[i] - list[j]
            if bisect_right(sum, diff) - bisect_left(sum, diff) != 0:
                print(list[i])
                return


func()

나다어

• 이진 탐색 함수는 정렬된 상태에서 실행해야 한다.
• 숫자를 감소시키는 반복문은 reversed 혹은 range(n, 0, -1) 처럼 간격을 -1로 주면 된다.
• for문을 종료하는 방법은 break /아니면 함수로 만들어서 return사용
• 배열에 특정 원소의 존재 여부를 찾으려면 bisect 라이브러리 사용 (in 연산자는 시간 오래 걸림)
• 해답 코드에서는 사용하지 않았지만,
  중복을 제거하기 위해 set 자료형을 사용할 수 있다.