Balls of Three Colors

• 문제는 설명은 심플한데, 풀이를 떠올리는 것이 쉽지는 않다
• 풀이 방법은 다양하게 있을 것 같다. 적어도 내가 푼 방법은 에디토리얼의 풀이와 다르긴 하다.
• 내 풀이 기준으로 쓰면, 우선 빨강과 초록를 배치하고, 그 사이사이에 파랑을 끼워넣는 것을 목표로 한다
  • 빨강과 초록 사이사이에 파랑을 배치할때, 파랑을 넣을 수 있는 곳은 r+g+1 개이다. 그리고 각각에는 파랑이 최대 1개씩만 들어갈수 있다.
  • 그리고 빨강-빨강 또는 초록-초록 사이에는 반드시 파랑이 들어가야 한다. 이런곳이 k개라고 하면, (r+g+1 - k)개의 빈 곳 중에서 파랑 (b-k)개가 들어갈 곳들을 고르는 경우의 수이므로 C(r+g+1-k,b-k) 개가 된다.
  • k를 결정하는 것은 빨강과 초록을 배치했을때, 같은 색끼리 연속된 run의 개수이다. 이것을 l이라고 하자
  • 공과 공 사이의 개수가 r+g-1 개인데, 이중에서 양쪽의 색깔이 다른 것의 개수는 l-1 개이므로, 색깔이 같은 것의 개수 k = (r+g-1) - (l-1) = r+g-l 이다
  • 이제 run의 개수가 l이 되도록 빨강과 초록를 배치하는 방법의 수를 구하자.
  • l=2n 이면, 빨강을 n그룹으로 나누고, 초록도 n그룹으로 나눈 다음에 교차하면서 배치해주면 된다.
    • 가짓수는 C(r-1,n-1) * C(g-1,n-1) * 2 가 된다. 2를 곱하는 것은 빨-초-빨-초.. 배치와 초-빨-초-빨.. 배치 두가지가 가능하기 때문
  • l=2n+1이면 빨강 n+1개와 초록 n개로 나누거나 빨강 n개와 초록 n+1개로 나누어서 교차시킨다.
    • 가짓수는 C(r-1,n) * C(g-1,n-1) + C(r-1,n-1) * C(g-1,n) 이다
  • 이제 이것을 합쳐서 계산하면 된다.
    • sum_{i=2,3,…} {l=i 가 되도록 빨강,초록을 배치하는 방법} * {l=i일때 파랑을 배치하는 방법} 을 구하면 된다.
    • O(r+g+b)의 전처리 이후에 C(x,y)를 모두 O(1)에 계산하도록 하면, 총 시간복잡도는 O(r+g+b) 가 된다.
• 에디토리얼에서는 별해로 f(a,b,c)=f(a-1,b-1,c)+f(a-1,b,c-1)+f(a,b-1,c-1)+2f(a-1,b-1,c-1) 이라는 점화식을 이용해서 O(r+g+b) 에 푸는 것도 가능하다고 간단히 언급하고 있다.
  • 하지만.. 어떻게 저런 점화식이 유도되는지도, 저 점화식을 어떻게 선형에 계산하는지도 둘다 모르겠다;

"""Solution code for "BOJ 32914. Balls of Three Colors".

- Problem link: https://www.acmicpc.net/problem/32914
- Solution link: http://www.teferi.net/ps/problems/boj/32914

Tags: [combinatoric]
"""

from teflib import modcomb


MOD = 998_244_353


def main():
    modcomb.set_mod(MOD)

    r, g, b = [int(x) for x in input().split()]

    answer = 0
    for i in range(max(2, r + g - b), (min(r, g) + 1) * 2):
        b_count = modcomb.comb(i + 1, b - (r + g - i))
        ii = i // 2 - 1
        if i % 2 == 0:
            rg_count = modcomb.comb(r - 1, ii) * modcomb.comb(g - 1, ii) * 2
        else:
            rg_count = (
                modcomb.comb(r - 1, ii) * modcomb.comb(g - 1, ii + 1)
            ) + modcomb.comb(r - 1, ii + 1) * modcomb.comb(g - 1, ii)
        answer += rg_count * b_count

    print(answer % MOD)


if __name__ == '__main__':
    main()

• Dependency: teflib.mobcomb.comb