김낭만

재밌거나 필요하다고 생각하는 걸 만듭니다.

백준(BOJ) 1234 크리스마스 트리 파이썬(Python)

문제 링크

간단한 DP 문제.
높이가 n이고, 남은 장난감이 각각 a, b, c개인 문제는 높이가 n - 1인 트리에서 남은 장난감이 a, b, c에서 총 n개 줄어든 부분문제로 구성된다.
이때, 장난감이 줄어드는 방법의 가짓수는 최대 7가지(1개 3가지, 2개 3가지, 3개 1가지)이므로 가능한 경우에 대해 완전탐색을 하되, DP 문제이므로 메모이제이션을 해주면 쉽게 풀 수 있다.
solve(a, b, c, i)는 장난감이 각각 a, b, c개 있고 레벨이 i일 때의 답이다.
이때, 공의 개수가 같다면 색깔이 달라도 경우의 수 개수는 같으므로(즉, solve(1, 2, 3, 2)나 solve(3, 2, 1, 2)나 같으므로) a, b, c를 정렬해서 캐시 적중률을 올릴 수 있다.
이때, 장난감 i개를 골랐다고 끝이 아니라, i개를 배열하는 경우의 수를 구해야 한다.
모든 장난감이 같은 색이라면 경우의 수는 1가지이다.
2가지 색의 장난감을 배열하는 경우 i개의 자리 중 1번 색을 넣을 i // 2개의 자리를 구하고 나머지에는 2번 색을 넣으면 된다(즉, combinations(i, i // 2)).
3가지 색의 장난감을 배열하는 경우 i개의 자리 중 1번 색을 넣을 i // 3개의 자리를 구하고, 나머지 i - i // 3개의 자리 중 i // 3개 자리에 2번 색을 넣고, 나머지 자리에 3번 색을 넣으면 된다(즉, combinations(i, i // 3) * combinations(i - i // 3, i // 3)).
이러한 경우들만 잘 계산하면 큰 어려움 없이 답을 구할 수 있다.
여담으로 탑다운 방식의 DP를 쓰기 용이한 문제이므로 functools.cache를 사용하면 쉽게 구현할 수 있다.

#!python

from functools import cache, reduce

n, r, g, b = map(int, input().split())

# aCb를 계산하는 함수
@cache
def combinations(a, b):
    return reduce(lambda x, y: x * y, range(b + 1, a + 1), 1) // reduce(lambda x, y: x * y, range(1, a - b + 1), 1)

# r, g, b를 정렬해서 a, b, c에 넣고 i는 레벨 k를 의미
# 트리를 아래에서부터 채워나감
@cache
def solve(a, b, c, i):

    # 레벨 0까지 왔다 = 성공이므로 경우의 수 +1
    if not i:
        return 1
    
    ret = 0

    # 하나의 장난감만 사용하는 경우
    if a >= i:
        ret += solve(a - i, b, c, i - 1)
    if b >= i:
        ret += solve(*sorted((a, b - i, c)), i - 1)
    if c >= i:
        ret += solve(*sorted((a, b, c - i)), i - 1)

    # 2종류의 장난감을 사용하는 경우
    # a >= b >= c이므로 a >= i // 2면 b와 c도 그렇고, b >= i // 2면 c도 그렇다
    if i % 2 == 0:
        if a >= i // 2:
            ret += solve(a - i // 2, b - i // 2, c, i - 1) * combinations(i, i // 2)
            ret += solve(*sorted((a - i // 2, b, c - i // 2)), i - 1) * combinations(i, i // 2)
        if b >= i // 2:
            ret += solve(*sorted((a, b - i // 2, c - i // 2)), i - 1) * combinations(i, i // 2)

    # 3종류의 장난감을 사용하는 경우
    # 마찬가지로 a >= b >= c이므로 a만 확인하면 됨
    if i % 3 == 0 and a >= i // 3:
        ret += solve(a - i // 3, b - i // 3, c - i // 3, i - 1) * combinations(i, i // 3) * combinations(i - i // 3, i // 3)

    return ret

# r, g, b를 정렬해서 넣고 아래쪽부터 채우므로 n을 넣는다
print(solve(*sorted((r, g, b)), n))

2024

맨 위로 이동 ↑

2023

세그먼트 트리

개요 선형적인 자료구조에서는 값에 접근하는 데에 \(O(1)\)이면 충분하지만, 대신 부분합을 구하는 데에는 \(O(N)\)이 필요하다. 그렇다면 이 자료구조를 이진 트리로 구성하면 어떨까? 값에 접근하는 데에 걸리는 시간이 \(O(\lg N)\)으로 늘어나지만 대신 부분합을 구하...

벨만-포드 알고리즘

개요 다익스트라 알고리즘과 함께 Single Sourse Shortest Path(SSSP) 문제를 푸는 알고리즘이다. 즉, 한 노드에서 다른 모든 노드로 가는 최단 경로를 구하는 알고리즘이다. 다익스트라 알고리즘보다 느리지만, 음수 가중치 간선이 있어도 작동하며, 음수 가중치 사...

다익스트라 알고리즘

개요 다익스트라 알고리즘은 Single Sourse Shortest Path(SSSP) 문제를 푸는 알고리즘 중 하나이다. 즉, 한 노드에서 다른 모든 노드로 가는 최단 경로를 구하는 알고리즘이다. 단, 다익스트라 알고리즘은 음수 가중치 엣지를 허용하지 않는다. 이 경우에는 벨만-...

맨 위로 이동 ↑