[BOJ] 1301 비즈 공예

1301 비즈 공예 문제 링크
문제 풀이 코드 링크

문제 요약

구슬들을 이용해 목걸이를 만든다. N개의 종류의 구슬을 이용하는데, 연속된 3개의 구슬은 모두 색이 달라야한다. 단, 각 구슬의 개수는 제한되어있다. 이 때 목걸이를 만들 수 있는 경우의 수를 출력한다.

아이디어

혼자 힘으로 못 풀고, 또 다시 구글링의 힘을... :)
DP 문제이다. 매 단계의 선택에 영향을 주는 요소는 (1)남은 구슬의 개수와 (2)전과 그 전에 꿴 구슬의 종류이다.
따라서 구슬의 종류(최대 5종류)+이전 단계의 구슬의 종류에 따라 7차원 배열(...)로 설계.
[0번 구슬의 남은 수][1번 구슬의 남은 수]…[4번 구슬의 남은 수][1단계 전의 구슬의 종류][2단계 전의 구슬의 종류]

TOP-DOWN 형식으로 구슬이 모두 남아있을 때부터 차례대로 DP 배열을 갱신한다. 매 단계에서 다음 재귀함수를 호출할 때에는 현재 꿴 구슬의 개수를 빼고 전달한다. 함수에서는 남은 구슬로 만들 수 있는 경우의 수를 계산하고, 해당 계산 결과를 DP 배열에 저장한다. 같은 경우에 대해 해당 함수가 호출되면, 해당 DP 값을 반환한다.

구현

import sys

input = sys.stdin.readline

# Get input
N = int(input())

beedsNum = [0 for i in range(5)]
for i in range(N) :
  beedsNum[i] = int(input())
totalNum = sum(beedsNum)

# pre2와 pre1은 각각 2단계, 1단계전 구슬의 종류
# 구슬의 남은 개수 배열 beedsNum[]은 전역 변수로 매번 갱신됨.
def addBeeds(cur, pre2, pre1) :
  # 모든 구슬을 다 사용했다면 목걸이를 끝까지 완성하였으므로 가능한 경우의 수. 따라서 1을 반환.
  if cur >= totalNum :
    return 1
  
  # DP 배열에 값이 저장되어있으면 해당 값을 반환.
  if dp[beedsNum[0]][beedsNum[1]][beedsNum[2]][beedsNum[3]][beedsNum[4]][pre2][pre1] != -1 :
    return dp[beedsNum[0]][beedsNum[1]][beedsNum[2]][beedsNum[3]][beedsNum[4]][pre2][pre1]

  sum = 0  
  for i in range(N) :
    if i != pre1 and i != pre2 and beedsNum[i] > 0:
      beedsNum[i] -= 1
      sum += addBeeds(cur + 1, pre1, i)
      beedsNum[i] += 1

  dp[beedsNum[0]][beedsNum[1]][beedsNum[2]][beedsNum[3]][beedsNum[4]][pre2][pre1] = sum
  return sum

# DP 배열 선언(...)
dp = [[[[[[[-1 for i in range(N + 1)] for j in range(N + 1)] for e in range(11)] for d in range(11)] for c in range(11)] for b in range(11)] for a in range(11)]
print(addBeeds(0, -1, -1))