[BOJ] 백준 2632번 문제 풀이 (Python)

1. 문제 풀이

피자 A, 피자 B에서 나올 수 있는 모든 경우의 수를 dictionary에 저장하였다.

• case1: 피자 A에서 나올 수 있는 피자 조각의 경우의 수
  • 예를 들어, 합이 5인 피자 조각이 5개라면 case1[5] = 5 가 된다.
• case2: 피자 B에서 나올 수 있는 피자 조각의 경우의 수
  • 예를 들어, 합이 4인 피자 조각이 3개라면 case2[4] = 3 가 된다.

하나의 피자에서 나올 수 있는 모든 경우의 수 탐색하기

• 하나의 피자에 존재하는 조각의 갯수가 1000개 이하이므로 반복문 2개를 중첩하여 모든 경우의 수를 탐색할 수 있다.

피자 조각: 1 2 3 4 5

(조각     ⇒ 합)
1        ⇒ 1
1 2      ⇒ 3
1 2 3    ⇒ 6 
1 2 3 4  ⇒ 10

2       ⇒ 2
2 3     ⇒ 5 
2 3 4   ⇒ 9
2 3 4 5 ⇒ 14

3       ⇒ 3
3 4     ⇒ 7 
3 4 5   ⇒ 12
3 4 5 1 ⇒ 13

4       ⇒ 4
4 5     ⇒ 9
4 5 1   ⇒ 10
4 5 1 2 ⇒ 12

5       ⇒ 5
5 1     ⇒ 6
5 1 2   ⇒ 8
5 1 2 3 ⇒ 11

1 2 3 4 5 ⇒ 15

합이 K인 경우의 수 찾기

• case1에 존재하는 피자 조각 중 합이 S인 것이 존재한다면 case2에 합이 K-S인 것이 존재하는 지 체크하면 된다. case1[S] * case2[K-S]

2. 코드

import sys
from collections import defaultdict
def input(): return sys.stdin.readline().rstrip()

def find_case(pizza, length):
    case = defaultdict(int)
    for i in range(length):
        temp = pizza[i:] + pizza[:i]
        pre = 0
        for num in temp:
            pre += num
            case[pre] += 1
    case[sum(pizza)] = 1
    return case

k = int(input())
n, m = map(int,input().split())
pizza_a = [int(input()) for _ in range(n)]
pizza_b = [int(input()) for _ in range(m)]

case1 = find_case(pizza_a, n)
case2 = find_case(pizza_b, m)

result = case1.get(k, 0) + case2.get(k, 0)
for num in case1:
    if k-num in case2:
        result += case1[num] * case2[k-num]
print(result)