[백준/1234] 크리스마스 트리 (Python)

Problem : https://www.acmicpc.net/problem/1234

1234번: 크리스마스 트리

 

Difficulty : Gold 2

 

Status : Solved

 

Time : 00:35:42

 

---

 

문제 설명

 

 

오민식은 오늘이 크리스마스라고 생각해서, 크리스마스 트리를 만들려고 한다. 트리는 N개의 레벨로 이루어져 있다. 위에서부터 레벨1, ... 레벨 N이다. 또, 민식이는 빨강, 파랑, 초록색의 장난감을 가지고 있다. 그리고 민식이는 이 장난감을 일정한 규칙에 의해서 장식하려고 한다.

레벨 K에는 딱 K개의 장난감이 있어야 한다. 또, 각 레벨에 놓으려고 선택한 색이 있으면, 그 색의 장난감의 수는 서로 같아야 한다. 예를 들어, 레벨 3에 장난감을 놓으려고 할 때, 빨강 2, 파랑 1과 같이 놓으면, 빨강과 파랑의 수가 다르기 때문에 안 된다. 하지만, 레벨 4에 빨강 2, 파랑 2와 같이 놓으면, 가능하다.

N과, 장난감의 수가 주어질 때, 트리를 장식하는 경우의 수를 출력하는 프로그램을 작성하시오.

 

 

입력 및 출력

 

입력

 

첫째 줄에 트리의 크기 N, 빨강의 개수, 초록의 개수, 파랑의 개수가 주어진다. N은 10보다 작거나 같다. 빨강, 초록, 파랑의 개수는 0보다 크거나 같고, 100보다 작거나 같다.

 

출력

 

첫째 줄에 경우의 수를 출력한다. 만약 주어진 장난감으로 트리를 장식할 수 없으면 0을 출력한다. 문제의 정답은 263-1보다 작거나 같다.

 

입력 예시

 

2 1 1 1

 

출력 예시

 

6

 

 

---

 

풀이

 

 

풀이는 이후에 작성한다 (여백이 작아서 나중에 적도록 한다 ^^)

 

풀이 코드

N, R, G, B = map(int, input().split())

def comb(n, r):
  if r == n or r == 0 :
    return 1
  return comb(n-1, r-1) + comb(n-1, r)

dp = [[[[0]*(B+1) for _ in range(G+1)] for _ in range(R+1)] for _ in range(N+1)]

if B :
  dp[1][R][G][B-1] = 1
if G :
  dp[1][R][G-1][B] = 1
if R :
  dp[1][R-1][G][B] = 1
  
for i in range(1, N) :
  for r in range(R+1) :
    for g in range(G+1) :
      for b in range(B+1) :
        if not dp[i][r][g][b] :
          continue
        d, t = (i+1) // 2, (i+1) // 3
        if (i+1) % 3 == 0 and r >= t and g >= t and b >= t :
          dp[i+1][r-t][g-t][b-t] += dp[i][r][g][b] * comb(3*t, t) * comb(2*t, t)
        if (i+1) % 2 == 0 :
          if r >= d and g >= d :
            dp[i+1][r-d][g-d][b] += dp[i][r][g][b] * comb(2*d, d)
          if g >= d and b >= d :
            dp[i+1][r][g-d][b-d] += dp[i][r][g][b] * comb(2*d, d)
          if r >= d and b >= d :
            dp[i+1][r-d][g][b-d] += dp[i][r][g][b] * comb(2*d, d)
        if r > i :
          dp[i+1][r-i-1][g][b] += dp[i][r][g][b]
        if g > i :
          dp[i+1][r][g-i-1][b] += dp[i][r][g][b]
        if b > i :
          dp[i+1][r][g][b-i-1] += dp[i][r][g][b]

ans = 0
for i in dp[-1] :
  for j in i :
    for k in j :
      ans += k
print(ans)

풀이 완료!ㅅ