문제

N(<= 5)종류의 서로 다른 색상의 구슬을 가지고 있다.
구슬을 임의의 연속된 3개의 구슬의 색이 모두 다 다르게 하려고 한다.
이 때, 만들 수 있는 목걸이의 경우의 수를 구한다.

아이디어

핵심 아이디어

• 메모이제이션을 통해 중복계산을 방지한다.

위의 그림에서 RGB 3 색상을 사용하고, val의 O는 임의의 색이라 가정할 때
N idx의 값을 채우기 위해서는 종류별 남은 구슬의 갯수들과 이전/이이전에 사용한 구슬의 색의 정보만 필요하다.
즉, 남은 구슬을 채우기 위해서는 N - 2, N - 1번째 구슬을 제외하고는 그 이전에 어떤 순서로 구슬을 채웠는지는 알 필요가 없다.

이 사실을 염두해두고 아래 상황을 본다위의 상황에서 6 idx이후의 모든 경우의 수를 구했다 가정하고, 아래 상황에 도착했을 때를 생각한다.
6 idx이전의 값이 약간 다르지만 구슬종류별 사용갯수와 남은 갯수는 동일하다.
이 때, 6 idx에 B를 채웠을 경우 이후의 상황은 이미 구했기 때문에 다시 반복할 필요가 없다.
6 idx 이후의 구슬 경우의 수는 남은 갯수와 4, 5 idx의 구슬 순서만이 결정하기 때문에 1번 상황과 2번 상황 모두 같다.

구현 아이디어

• 7 차원의 메모이제이션 배열을 사용한다.
  • 1,2,3,4,5 차원은 구슬 종류별 남은 갯수를 의미한다.
  • 6,7 차원은 이전/ 이이전 사용 구슬 종류를 의미한다.

코드

import sys

input = sys.stdin.readline

N = int(input())
beads = [0 for _ in range(5)]
for n in range(N):
    beads[n] = int(input())
memo = [[[[[[[-1 for _ in range(6)] for _ in range(6)] for a in range(11)] for b in range(11)] for c in range(11)] for d in range(11)] for e in range(11)]
total = sum(beads)


def memoization(size: int, pp: int, p: int) -> int:
    if size == total:
        return 1

    if memo[beads[0]][beads[1]][beads[2]][beads[3]][beads[4]][pp][p] != -1:
        return memo[beads[0]][beads[1]][beads[2]][beads[3]][beads[4]][pp][p]
    # 한번 확인한 경우, 다시 확인할 필요가 없다.
    ret = 0

    for i, b in enumerate(beads):
        if i == p or i == pp: continue
        if b == 0: continue
        beads[i] -= 1
        ret += memoization(size + 1, p, i)
        beads[i] += 1

    memo[beads[0]][beads[1]][beads[2]][beads[3]][beads[4]][pp][p] = ret
    return ret

print(memoization(0, -1, -1))