백준: 14238 출근 기록

문제

14238번: 출근 기록

 

문제 풀이

 세 직원 A, B, C가 있다. 회사에는 특별한 룰을 가지고 있어서 하루에 한 명만 출근한다. 각 직원에 따라 출근할 수 있는 규칙이 정해져 있다. A의 경우 매일 출근할 수 있으며, B의 출근한 다음날은 쉬어야 한다. C의 경우 출근한 다음날과 그다음 날 모두 쉬어야 한다는 규칙이 있다.

 

 따라서 dp의 배열을 5차원 배열로 `dp[a][b][c][전전날][전날]`과 같이 기록하여, 중복으로 경우의 수를 탐색하지 않도록 하여야 한다. 또한, 사용 가능한 A, B, C의 수에 따라 경우의 수를 찾아가며 전날에 B가 일하였거나, 전과 전전날에 C가 일한 경우는 경우의 수에 추가되지 않도록 하여 원하는 답을 찾을 수 있다.

 

코드

from sys import stdin


def dfs(a, b, c, prev):

    # a, b, c의 개수가 초기 카운트와 같으면 찾은 것
    if [a, b, c] == cnt:
        print(''.join(answer))
        exit(0)

    if dp[a][b][c][prev[0]][prev[1]]:
        return False

    dp[a][b][c][prev[0]][prev[1]] = True

    # A개수 만큼 사용하지 않았다면
    if a + 1 <= cnt[A]:
        answer[a + b + c] = 'A'
        if dfs(a + 1, b, c, [prev[1], A]):
            return True
    # B개수 만큼 사용하지 않았다면
    if b + 1 <= cnt[B]:
        answer[a + b + c] = 'B'
        # 전날에 선택한 것이 B가 아니라면
        if prev[1] != B:
            if dfs(a, b + 1, c, [prev[1], B]):
                return True
    # C개수 만큼 사용하지 않았다면
    if c + 1 <= cnt[C]:
        answer[a + b + c] = 'C'
        # 전전날과 전날에 선택한 것이 C가 아니라면
        if prev[0] != C and prev[1] != C:
            if dfs(a, b, c + 1, [prev[1], C]):
                return True
    return False


if __name__ == '__main__':
    A, B, C = 0, 1, 2
    s = list(stdin.readline().rstrip())
    length = len(s)
    answer = [''] * length
    # a, b, c 개수 카운트.
    cnt = [s.count(word) for word in ('A', 'B', 'C')]
    dp = [[[[[False for _ in range(3)] for _ in range(3)] for _ in range(length)] for _ in range(length)] for _ in range(length)]
    dfs(0, 0, 0, [0, 0])
    print(-1)

 5차원 배열을 생성하는 DP문제를 처음 풀어보았다. 풀다 보면서 느끼는건 DP를 풀기 위해 적절한 가지치기를 설정하고, 중복으로 탐색될 경우를 방지하기 위한 메모이제이션을 잘 설계하는 것이 중요한 것 같다.