[BOJ] 14238: 출근 기록 (Python)

문제 링크

https://www.acmicpc.net/problem/14238

풀이 전략

Dynamic Programming

5차원 배열을 이용한 memoization을 이용해 답을 구할 수 있다.

dp의 배열을 5차원 배열로 dp[a][b][c][전전날][전날]과 같이 기록하여, 중복으로 경우의 수를 탐색하지 않도록 하여야 한다.

또한, 사용 가능한 A, B, C의 수에 따라 경우의 수를 찾아가며 전날에 B가 일하였거나, 전과 전전날에 C가 일한 경우는 경우의 수에 추가되지 않도록 하여 답을 구할 수 있다.

구현

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)

삽질

모든 순열을 전부 구해 리스트에 저장하고, 하나씩 꺼내서 확인해보는 방식을 이용했는데, 당연하게도 메모리 초과가 발생했다.

from itertools import permutations

if __name__ == "__main__":
    S = input()
    length = len(S)
    words = list((map(''.join, permutations(S))))

    for word in words:
        dp_c = [0] * (length+2)
        dp_b = [0] * (length+2)
        for i in range(length):
            if (s:=word[i]) == 'B':
                if dp_b[i] != 1:
                    dp_b[i+1] = 1
                else:
                    break
            elif s == 'C':
                if dp_c[i] != 1:
                    dp_c[i+1] = 1
                    dp_c[i+2] = 1
                else:
                    break
        else:
            print(word)
            break
    else:
        print(-1)

이 외에도 엄청난 삽질을 했다. 덕분에 차원을 많이 다루는 문제들은 쉽게 풀 수 있게 되었다.