문제 정보
문제 정리
크리스마스 트리의 i층에는 i개의 칸이 있으며, 각 칸에는 색깔을 배치할 수 있다. 또한, 어떤 층에 놓여 있는 색깔은 그 수가 같아야 한다. 트리의 크기와 빨강, 파랑, 초록의 개수가 주어질 때, 트리를 장식하는 경우의 수를 구하는 문제이다.
풀이
우선 문제를 파악해보자.
트리의 1층에는 1칸, 2층에는 2칸, ... N층에는 N칸이 있다. (a, b, c)를 현재 층에 놓인 빨강, 초록, 파랑의 수라고 둔 다음 가능한 케이스를 생각해보자. N <= 10이므로 10층까지의 케이스를 전부 나열해보면,
1층 : (1, 0, 0)
2층 : (2, 0, 0), (1, 1, 0)
3층 : (3, 0, 0), (1, 1, 1)
4층 : (4, 0, 0), (2, 2, 0)
5층 : (5, 0, 0)
6층 : (6, 0, 0), (3, 3, 0), (2, 2, 2)
7층 : (7, 0, 0)
8층 : (8, 0, 0), (4, 4, 0)
9층 : (9, 0, 0), (3, 3, 3)
10층 : (10, 0, 0), (5, 5, 0)
(5, 0, 0), (0, 5, 0), (0, 0, 5)와 같이 순서를 바꾸는 경우는 너무 길어질 것 같아서 하나의 케이스로 퉁쳤다.
여기서 알 수 있는 것은 현재 층이 i일 때,
- (i, 0, 0)은 모든 층에 들어간다는 것.
- i % 2 == 0일 경우 (i/2, i/2, 0)이 가능하다는 것.
- i % 3 == 0일 경우 (i/3, i/3, i/3)이 가능하다는 것.
이렇게 보니 DP의 냄새가 난다.
이전 층까지 사용한 색깔의 수를 알고 있는 경우, 현재 층에 사용할 색깔을 계산할 수 있으며, 이걸 부분 문제로 볼 경우 아래처럼 점화식을 세울 수 있다.
DP[i][j][k][l] = (1~i층을 꾸밀 때 사용한 빨강이 j, 초록이 k, 파랑이 l일 때의 경우의 수).
점화식을 이렇게 세워두긴 했지만, j, k가 정해질 경우 l은 자동 고정되고, DP[i]를 계산하기 위해선 DP[i-1]의 값만 필요하므로, 한 층마다 새 배열을 만들 경우 사실상 2차원 배열만으로 계산할 수 있다.
문제 태그의 조합론은 그냥 한 층에 빨간색 m개, 파란색 n개를 배치하는 경우의 수? 정도의 계산만 필요하다. N이 10까지밖에 안되서 조합 값을 직접 하드코딩했다.
초기화 과정에서 DP[1] 값을 저장해야 하는데 색이 0개일 경우 DP[1]의 값이 정의되지 않는데 그걸 감안하지 않아서 런타임 에러가 나왔다. 그거랑은 별개로 구현이 좀 빡센 것 같다.
코드
# 백준 1234
import io
input = io.BufferedReader(io.FileIO(0), 1<<18).readline
C = [[], [], [2], [0, 6], [6], [], [20, 90], [], [70], [0, 1680], [252]]
def solve(N, R, G, B):
# 색이 부족한 경우
if R+G+B < N*(N+1)//2:
return 0
# DP 초기값 처리
prevDP = [[0 for _ in range(G+1)] for _ in range(R+1)]
if B > 0:
prevDP[0][0] = 1
if G > 0:
prevDP[0][1] = 1
if R > 0:
prevDP[1][0] = 1
for lv in range(2, N+1):
accSum = lv*(lv+1)//2
DP = [[0 for _ in range(G+1)] for _ in range(R+1)]
# 현재 층까지 사용한 빨강의 총 개수가 totalR, 초록의 총 개수가 totalG일 때, DP 계산
for totalR in range(R+1):
for totalG in range(G+1):
totalB = accSum - totalR - totalG
if totalB < 0 or totalB > B:
continue
# 현재 층에서 사용할 색 조합 저장
caseL = [[lv, 0, 0, 1], [0, lv, 0, 1], [0, 0, lv, 1]]
if lv % 2 == 0:
t = lv//2
caseL.append([t, t, 0, C[lv][0]])
caseL.append([t, 0, t, C[lv][0]])
caseL.append([0, t, t, C[lv][0]])
if lv % 3 == 0:
t = lv//3
caseL.append([t, t, t, C[lv][1]])
# 색 조합을 바탕으로 DP 계산
for useR, useG, useB, value in caseL:
prevR = totalR - useR
prevG = totalG - useG
prevB = totalB - useB
if 0 <= prevR <= R and 0 <= prevG <= G and 0 <= prevB <= B:
DP[totalR][totalG] += prevDP[prevR][prevG] * value
prevDP = DP
# 사용한 색의 개수가 R, G, B인 값들의 DP 총합 구하기.
result = 0
curUse = N*(N+1)//2
for y in range(R+1):
for x in range(G+1):
curB = curUse - y - x
if 0 <= curB <= B:
result += prevDP[y][x]
return result
def main():
N, R, G, B = map(int, input().split())
print(solve(N, R, G, B))
main()
