문제

길이가 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이 되어 최대가 된다.

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

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

www.acmicpc.net/problem/1744

  • 입력

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

  • 출력

수를 합이 최대가 나오게 묶었을 때 합을 출력한다. 정답은 항상 231보다 작다.

과정

수를 묶을 때 위치에 제한이 없기 때문에 정렬한 다음에 간단한 조건에 따라 수를 묶으면 된다고 생각했다. 조건은 아래와 같다.

  1. n1 and n2 > 1: 큰 수부터 차례대로 묶어서 곱한다.
  2. n1 > 1 and n2 == 1: 곱한 값보다 더한 값이 크기 때문에 묶지 않고 더한다.
  3. n1 and n2 <= 0: 작은 수부터 차례대로 묶어서 곱한다. 하나의 수가 0이더라도 곱한 값이 음수보단 크므로 똑같이 곱한다.

아래는 이 조건에 맞춰 구현했지만 실패한 코드이다. 복잡하고 구린 코드이니 너무 자세히 보진 말자

import sys

n = int(sys.stdin.readline())
arr = [int(sys.stdin.readline()) for _ in range(n)]
# 배열을 오름차순으로 정렬
arr = sorted(arr)
res = 0

# 숫자 <= 0 인 경우
i = 0
while True:
    # index가 범위 밖이거나, 첫 번째 숫자가 음수가 아니거나, 두 번째 숫자가 0인 경우
    if i >= n or arr[i] >= 0 or arr[i+1] == 0:
        break
    # 마지막  index이거나, 두 번째 숫자가 양수인 경우
    elif (i == n - 1 and arr[i] < 0) or arr[i+1] > 0:
        res += arr[i]
        break
    # 둘 다 음수인 경우
    else:
        res += arr[i] * arr[i+1]
        i += 2

# 숫자 > 0 인 경우
i = 0
# 내림차순으로 정렬
arr = sorted(arr, reverse=True)
while True:
    # index가 범위 밖이거나 첫 번째 숫자가 양수가 아닌 경우
    if i >= n or arr[i] <= 0:
        break
    # 마지막 숫자이거나 두 번째 숫자가 양수가 아닌 경우
    elif i == n - 1 or arr[i+1] <= 0:
        res += arr[i]
        break
    # 첫 번째 숫자가 1인 경우
    if arr[i] == 1:
        res += arr[i]
        i += 1
    # 두 번째 숫자가 1인 경우
    elif arr[i+1] == 1:
        res += arr[i] + 1
        i += 2
    # 둘 다 1보다 큰 양수인 경우
    else:
        res += arr[i] * arr[i+1]
        i += 2

print(res)

보다시피 복잡한데 인덱스 에러까지 발생해서 실패한 코드이다. 복잡하게 조건문으로 조건 체크도 열심히 한 것 같은데 어디서 인덱스 에러가 발생하는지 열심히 보니까 첫 번째 while loop의 i >= n or arr[i] >= 0 or arr[i+1] == 0 이런 경우나 두 번째 while loop의 i == n - 1 or arr[i+1] <= 0 이런 경우에서 i가 이미 index 밖인데 arr[i] 혹은 arr[i+1]에 접근하는 부분이 하나의 조건문에 포함되어 있어서 인덱스 에러가 발생했다.

import sys

n = int(sys.stdin.readline())
arr = [int(sys.stdin.readline()) for _ in range(n)]
arr = sorted(arr)
res = 0

i = 0
while True:
    if i >= n:
        break
    elif i == n - 1 and arr[i] < 0:
        res += arr[i]
        break
    elif arr[i] >= 0 or arr[i+1] == 0:
        break
    elif arr[i+1] > 0:
        res += arr[i]
        break
    else:
        res += arr[i] * arr[i+1]
        i += 2


i = 0
arr = sorted(arr, reverse=True)
while True:
    if i >= n or arr[i] <= 0:
        break
    elif i == n - 1:
        res += arr[i]
        break
    elif arr[i+1] <= 0:
        res += arr[i]
        break
    if arr[i] == 1:
        res += arr[i]
        i += 1
    elif arr[i+1] == 1:
        res += arr[i] + 1
        i += 2
    else:
        res += arr[i] * arr[i+1]
        i += 2

print(res)

이렇게 고쳐서 기어이 통과를 받아냈지만 코드가 너무 구려서 찾아봤다.

최종코드

일단 index를 더 쉽게 다루기 위해서 n > 1인 경우와 n < 1인 경우를 arr1, arr2 배열에 나눠 담았다. 그리고 1은 어차피 곱하지 않고 더해야 하기 때문에 배열에 넣지 않고 바로 res 값에 더했다.

이렇게 배열 두 개로 나눠서 접근하니까 조건문도 훨씬 간단해져서 좋았다. 자세한 설명은 코드를 참고하자.

import sys

n = int(sys.stdin.readline())
arr1 = []
arr2 = []
res = 0

for _ in range(n):
    num = int(sys.stdin.readline())
    if num > 1:
        arr1.append(num)
    elif num < 1:
        arr2.append(num)
    # 1은 곱하지 않고 더해야 하므로 배열에 넣지 않고 바로 res에 더한다.
    else:
        res += 1

# arr1은 내림차순, arr2는 오름차순으로 정렬
arr1 = sorted(arr1, reverse=True)
arr2 = sorted(arr2)

for i in range(0, len(arr1), 2):
    # 마지막 원소인 경우 그냥 더해준다
    if i == len(arr1) - 1:
        res += arr1[i]
    else:
        res += arr1[i] * arr1[i+1]

for i in range(0, len(arr2), 2):
    # 마지막 원소인 경우 그냥 더해준다
    if i == len(arr2) - 1:
        res += arr2[i]
    else:
        res += arr2[i] * arr2[i+1]

print(res)

깔끔하게 코드 잘 짜고 싶다;_;

참고 사이트