[Python/파이썬] 백준 2473번 세 용액

2473번: 세 용액

문제

KOI 부설 과학연구소에서는 많은 종류의 산성 용액과 알칼리성 용액을 보유하고 있다. 각 용액에는 그 용액의 특성을 나타내는 하나의 정수가 주어져있다.  산성 용액의 특성값은 1부터 1,000,000,000까지의 양의 정수로 나타내고, 알칼리성 용액의 특성값은 -1부터 -1,000,000,000까지의 음의 정수로 나타낸다.

같은 양의 세 가지 용액을 혼합한 용액의 특성값은 혼합에 사용된 각 용액의 특성값의 합으로 정의한다. 이 연구소에서는 같은 양의 세 가지 용액을 혼합하여 특성값이 0에 가장 가까운 용액을 만들려고 한다. 

예를 들어, 주어진 용액들의 특성값이 [-2, 6, -97, -6, 98]인 경우에는 특성값이 -97와 -2인 용액과 특성값이 98인 용액을 혼합하면 특성값이 -1인 용액을 만들 수 있고, 이 용액이 특성값이 0에 가장 가까운 용액이다. 참고로, 세 종류의 알칼리성 용액만으로나 혹은 세 종류의 산성 용액만으로 특성값이 0에 가장 가까운 혼합 용액을 만드는 경우도 존재할 수 있다.

산성 용액과 알칼리성 용액이 주어졌을 때, 이 중 같은 양의 세 개의 서로 다른 용액을 혼합하여 특성값이 0에 가장 가까운 용액을 만들어내는 세 용액을 찾는 프로그램을 작성하시오.

 

입력

첫째 줄에는 전체 용액의 수 N이 입력된다. N은 3 이상 5,000 이하의 정수이다. 둘째 줄에는 용액의 특성값을 나타내는 N개의 정수가 빈칸을 사이에 두고 주어진다. 이 수들은 모두 -1,000,000,000 이상 1,000,000,000 이하이다. N개의 용액들의 특성값은 모두 다르고, 산성 용액만으로나 알칼리성 용액만으로 입력이 주어지는 경우도 있을 수 있다.

출력

첫째 줄에 특성값이 0에 가장 가까운 용액을 만들어내는 세 용액의 특성값을 출력한다. 출력해야하는 세 용액은 특성값의 오름차순으로 출력한다. 특성값이 0에 가장 가까운 용액을 만들어내는 경우가 두 개 이상일 경우에는 그 중 아무것이나 하나를 출력한다.

코드

import heapq
from itertools import combinations

n = int(input())
arr = list(map(int, input().split()))

comb_lst = combinations(arr, 3)

q = []
for el in comb_lst:
    heapq.heappush(q, (abs(sum(el)), el))
    
s, res = heapq.heappop(q)
for i in sorted(list(res)):
    print(i, end=" ")

최대로 n이 5000 들어온다면 5000개 중 3개를 뽑는 경우의 수는 n! / ((n - m)! * m!) = 20,820,835,000으로 200억이 넘는다...당연히 이렇게 하면 메모리 초과가 난다.

n = int(input())
arr = sorted(list(map(int, input().split())))
print(arr)
a, b, c = 0, 1, 2

ans = []
ans_sum = int(1e9) * 3

while a < (n-2):
    while c < n:
        while b < c:
            s = abs(arr[a] + arr[b] + arr[c])
            if s < ans_sum:
                ans_sum = s
                ans = [arr[a], arr[b], arr[c]]
            b += 1
        c += 1
        b = 1
    a += 1
    b, c= (a+1), (a+2)
    
print(" ".join(map(str, ans)))

이건 또 시간초과...당연하다 시간복잡도 대략 n ^ 3 나오니까 n이 500이어야 1초에 통과 가능하다

그래서 중간에 while b < c 이 부분 안에 아래와 같은 코드를 넣어서 합이 이전보다 커지면 더이상 연산하지 않게 해서 시간을 줄여봤더니 그냥 틀렸다

if prev < s:
                break

n = int(input())
arr = sorted(list(map(int, input().split())))

a, b, c = 0, 1, 2

ans = []
ans_sum = int(1e9) * 3
prev = int(1e9) * 3

while a < (n-2):
    while c < n:
        while b < c:
            s = abs(arr[a] + arr[b] + arr[c])
            if prev < s:
                break
            if s < ans_sum:
                ans_sum = s
                ans = [arr[a], arr[b], arr[c]]
            b += 1
            prev = s
        c += 1
        b = 1
    a += 1
    b, c= (a+1), (a+2)
    
print(" ".join(map(str, ans)))

그래서 세 포인터 중 왼쪽 포인터는 고정시켜놓고 가운데와 오른쪽 포인터를 움직이는 코드로 바꾸었다. 오른쪽 포인터는 이전과 다르게 배열의 오른쪽 끝에서 출발한다.

n = int(input())
arr = sorted(list(map(int, input().split())))

ans = []
ans_sum = int(1e9) * 3

for a in range(n-2):    # 왼쪽 고정
    b = a + 1
    c = n - 1
    while b < c:
        s = arr[a] + arr[b] + arr[c]
        if abs(s) <= abs(ans_sum):  # 새로운 값이 0에 더 가깝다면
            ans = [arr[a], arr[b], arr[c]]
            ans_sum = s        
        if s == 0:
            print(" ".join(map(str, ans)))
            exit()
        elif s < 0:
            b += 1
        else:
            c -= 1
            
print(" ".join(map(str, ans)))

a, b, c가 각각 왼쪽, 가운데, 오른쪽 포인터라고 할 때 왼쪽 포인터는 고정시켜놓고 b와 c에 투포인터 알고리즘을 사용하여 문제를 풀었다.

1. 용액의 특성값의 합(s)의 절댓값이 현재 저장된 최소값의 절댓값보다 작다면 그 때의 용액들의 특성값을 저장하고 최소값을 업데이트해준다.
2. s == 0이라면 더이상 볼 필요가 없으니 저장된 특성값들을 출력하고 프로그램을 종료한다.
3. s < 0이라면 합을 증가시키기 위해 b를 1 증가시킨다.
4. s > 0이라면 합을 감소시키기 위해 c를 1 감소시킨다.
5. 위의 과정들을 a가 n-2보다 작은 동안 반복한다. b와 c는 a의 값이 바뀔 때마다 각각 a+1, n-1로 업데이트된다.

 

참고로 이 코드는 PyPy3으로만 통과가 가능하고 Python3로는 시간 초과가 발생한다. Python3로 통과하기 위해서는 iter를 사용하여 파이썬스럽게 풀어야한다고 하는데...어차피 로직은 똑같으니 그냥 넘어간다.