백준 16198 파이썬 - 에너지 모으기

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

16198번: 에너지 모으기

문제

N개의 에너지 구슬이 일렬로 놓여져 있고, 에너지 구슬을 이용해서 에너지를 모으려고 한다.

i번째 에너지 구슬의 무게는 Wi이고, 에너지를 모으는 방법은 다음과 같으며, 반복해서 사용할 수 있다.

1. 에너지 구슬 하나를 고른다. 고른 에너지 구슬의 번호를 x라고 한다. 단, 첫 번째와 마지막 에너지 구슬은 고를 수 없다.
2. x번째 에너지 구슬을 제거한다.
3. Wx-1 × Wx+1의 에너지를 모을 수 있다.
4. N을 1 감소시키고, 에너지 구슬을 1번부터 N번까지로 다시 번호를 매긴다. 번호는 첫 구슬이 1번, 다음 구슬이 2번, ... 과 같이 매겨야 한다.

N과 에너지 구슬의 무게가 주어졌을 때, 모을 수 있는 에너지 양의 최댓값을 구하는 프로그램을 작성하시오.

코드

import sys
n = int(sys.stdin.readline())
arr = list(map(int, sys.stdin.readline().split()))
result = 0
def search(ar, tot):
    global result
    if len(ar) == 3:
        tot += ar[0] * ar[-1]
        result = max(tot, result)
        return
    for i in range(1, len(ar) - 1):
        search(ar[:i] + ar[i + 1:], tot + ar[i - 1] * ar[i + 1])
search(arr, 0)
print(result)

풀이

search함수에서는 ar 리스트를 인덱싱으로 하나씩 없애고 그 전과 뒤의 값을 곱해 tot에 더한다. ar의 길이가 3이 되면 return하는 조건을 넣어준다. result에 값을 저장하며 가장 큰 값을 구해 출력한다.