[Python/파이썬] 백준 1744번 수 묶기

1744번: 수 묶기

 

문제

길이가 N인 수열이 주어졌을 때, 그 수열의 합을 구하려고 한다. 하지만, 그냥 그 수열의 합을 모두 더해서 구하는 것이 아니라, 수열의 두 수를 묶으려고 한다. 어떤 수를 묶으려고 할 때, 위치에 상관없이 묶을 수 있다. 하지만, 같은 위치에 있는 수(자기 자신)를 묶는 것은 불가능하다. 그리고 어떤 수를 묶게 되면, 수열의 합을 구할 때 묶은 수는 서로 곱한 후에 더한다.

예를 들면, 어떤 수열이 {0, 1, 2, 4, 3, 5}일 때, 그냥 이 수열의 합을 구하면 0+1+2+4+3+5 = 15이다. 하지만, 2와 3을 묶고, 4와 5를 묶게 되면, 0+1+(2*3)+(4*5) = 27이 되어 최대가 된다.

수열의 모든 수는 단 한번만 묶거나, 아니면 묶지 않아야한다.

수열이 주어졌을 때, 수열의 각 수를 적절히 묶었을 때, 그 합이 최대가 되게 하는 프로그램을 작성하시오.

입력

첫째 줄에 수열의 크기 N이 주어진다. N은 50보다 작은 자연수이다. 둘째 줄부터 N개의 줄에 수열의 각 수가 주어진다. 수열의 수는 -1,000보다 크거나 같고, 1,000보다 작거나 같은 정수이다.

출력

수를 합이 최대가 나오게 묶었을 때 합을 출력한다. 정답은 항상 $2^{31}$보다 작다.

코드

def get_max_sum(arr):
    res = 0
    prev = 1001
    for i in range(len(arr)):
        if prev == 1001:
            prev = arr[i]
        else:
            res += prev * arr[i]
            prev = 1001
    res += (0 if prev == 1001 else prev)
    return res


n = int(input())
pos = []
neg = []
ans = 0
for _ in range(n):
    tmp = int(input())
    if tmp > 1:
        pos.append(tmp)
    elif tmp == 1:
        ans += 1
    else:
        neg.append(tmp)

pos.sort(reverse=True)  # 내림차순 정렬
neg.sort()  # 오름차순 정렬

ans += (get_max_sum(pos)+get_max_sum(neg))
print(ans)

 

처음에는 양수들은 큰 수부터 2개씩 묶어서 곱하고 그 합을 구했다. 그리고 1은 그냥 더했을 때 합을 최대로 만들 수 있으므로 그냥 더했다. 그리고 0이 있다면 가장 작은 음수부터 묶어서 곱해서 0으로 만들고 남은 음수들은 그냥 더하면 된다고 생각하여 그렇게 코드를 작성했는데 답이 아니었다. 왜냐하면 아래와 같은 예제는 음수끼리 곱해야 합이 최대가 되기 때문이다.

 

그래서 아예 입력받은 수를 1보다 큰 수인 pos, 0 이하의 수들인 neg 두 배열로 나누고, 1은 어떠한 배열에도 넣지 않고 처음부터 ans에 더하도록 하였다.

 

그리고 각 배열을 pos는 내림차순, neg는 오름차순으로 정렬하여 두 수열 모두 앞에서부터 2개씩 묶어서 곱한 값들의 합을 구하는 함수 get_max_sum()을 만들어서 pos와 neg가 합이 최대가 나오도록 묶었을 때의 합을 구하도록 하였다.