[14888] 연산자 끼워넣기 - Python

/* 문제 */
N개의 수로 이루어진 수열 A1, A2, ..., AN이 주어진다. 또, 수와 수 사이에 끼워넣을 수 있는 N-1개의 연산자가 주어진다. 연산자는 덧셈(+), 뺄셈(-), 곱셈(×), 나눗셈(÷)으로만 이루어져 있다.
우리는 수와 수 사이에 연산자를 하나씩 넣어서, 수식을 하나 만들 수 있다. 이때, 주어진 수의 순서를 바꾸면 안 된다.
예를 들어, 6개의 수로 이루어진 수열이 1, 2, 3, 4, 5, 6이고, 주어진 연산자가 덧셈(+) 2개, 뺄셈(-) 1개, 곱셈(×) 1개, 나눗셈(÷) 1개인 경우에는 총 60가지의 식을 만들 수 있다. 예를 들어, 아래와 같은 식을 만들 수 있다.
    1+2+3-4×5÷6
    1÷2+3+4-5×6
    1+2÷3×4-5+6
    1÷2×3-4+5+6
식의 계산은 연산자 우선 순위를 무시하고 앞에서부터 진행해야 한다. 또, 나눗셈은 정수 나눗셈으로 몫만 취한다. 음수를 양수로 나눌 때는 C++14의 기준을 따른다. 즉, 양수로 바꾼 뒤 몫을 취하고, 그 몫을 음수로 바꾼 것과 같다. 이에 따라서, 위의 식 4개의 결과를 계산해보면 아래와 같다.
    1+2+3-4×5÷6 = 1
    1÷2+3+4-5×6 = 12
    1+2÷3×4-5+6 = 5
    1÷2×3-4+5+6 = 7
N개의 수와 N-1개의 연산자가 주어졌을 때, 만들 수 있는 식의 결과가 최대인 것과 최소인 것을 구하는 프로그램을 작성하시오.

처음 떠오른 방법은 itertools를 이용해 가능한 연산자 순열을 모두 생성한 뒤 계산하는 것이었다.

# 시간 초과 코드
from itertools import permutations

n = int(input())
A = list(map(int, input().split()))
cnt = list(map(int, input().split()))   # 각 연산자의 개수
max_res = -1000000000
min_res = 1000000000
# 0:+, 1:-, 2:*, 3://
operators = []
for i in range(4):
    for _ in range(cnt[i]):
        operators.append(i)

for o in permutations(operators):
    tem = A[0]
    for i in range(n-1):
        if o[i] == 0:
            tem += A[i+1]
        elif o[i] == 1:
            tem -= A[i+1]
        elif o[i] == 2:
            tem *= A[i+1]
        else:
            if tem < 0:
                tem = -((-tem) // A[i+1])
            else:
                tem //= A[i+1]
    if tem < min_res:
        min_res = tem
    if tem > max_res:
        max_res = tem

print(max_res)
print(min_res)

위 코드는 Pypy3으로 채점하면 정답이지만 Python3에서는 시간 초과가 떴다. 모든 순열을 생성하고 그 순열의 모든 연산자에 대해 비교 연산을 한 후 계산해야 하므로 시간이 오래 걸렸을 것이다.

 

효율적으로 문제를 풀기 위해 dfs를 사용했다.

# 정답 코드
def dfs(pl, mi, mul, div, idx, tem_res):
    global n, max_res, min_res, A
    # 탐색이 끝나면 최대, 최소 업데이트
    if idx == n:
        max_res = max(max_res, tem_res)
        min_res = min(min_res, tem_res)
        return
    
    if pl > 0:
        dfs(pl - 1, mi, mul, div, idx + 1, tem_res + A[idx])
    if mi > 0:
        dfs(pl, mi - 1, mul, div, idx + 1, tem_res - A[idx])
    if mul > 0:
        dfs(pl, mi, mul - 1, div, idx + 1, tem_res * A[idx])
    if div > 0:
        dfs(pl, mi, mul, div - 1, idx + 1, int(tem_res / A[idx]))

n = int(input())
A = list(map(int, input().split()))
pl, mi, mul, div = map(int, input().split())   # 각 연산자의 개수
max_res = -1000000000
min_res = 1000000000
dfs(pl, mi, mul, div, 1, A[0])
print(max_res)
print(min_res)

순열을 생성하고 일일이 계산하는 것이 아니라, 남은 연산자의 개수를 dfs 함수에 넘겨주고 모든 연산자를 사용할 때까지 재귀적으로 dfs를 호출한다.

문제 분류에 백트래킹이라고 되어 있지만 모든 경우의 수를 탐색해야 하기 때문에 중간에 가지치기 코드를 넣지 않았다.

그리고 처음 코드에서는 음수의 나눗셈이 뜻대로 되지 않아 조건문을 추가했었는데, // 연산자 대신 int를 사용하도록 고쳤다.