[백준 2632번] 파이썬 - 피자판매

백준 2632 - 피자판매

시간 제한 2초, 메모리 제한 128MB

# 조건

• 고객이 두 종류의 피자 A와 B를 취급하는 피자가게에서 피자를 주문하고자 한다.
• <그림 1>과 같이 각 종류의 피자는 다양한 크기의 여러 개의 피자조각으로 나누어져 있다.
  • 각 조각에 쓰여진 숫자는 피자조각의 크기를 나타낸다.

• 고객이 원하는 피자의 크기를 이야기하면, 피자가게에서는 한 종류의 피자를 2 조각 이상 판매할 때는 반드시 연속된 조각들을 잘라서 판매한다.
• 이때 판매한 피자조각의 크기 합이 주문한 크기가 되어야 한다.
• 판매한 피자조각은 모두 A종류이거나, 모두 B종류이거나, 또는 A와 B 종류가 혼합될 수 있다.
  • 예를 들어서, <그림 1> 과 같이 잘라진 피자가 있을 때, 손님이 전체 크기가 7 인 피자를 주문하면, 피자 가게에서는 <그림2>와 같이 5 가지 방법으로 피자를 판매할 수 있다.

• 피자가게에서 손님이 원하는 크기의 피자를 판매하는 모든 방법의 가지 수를 계산하는 프로그램을 작성하시오

입력

• 첫 번째 줄에는 손님이 구매하고자 하는 피자크기를 나타내는 2,000,000 이하의 자연수가 주어진다.
• 두 번째 줄에는 A, B 피자의 피자조각의 개수를 나타내 는 정수 m, n 이 차례로 주어진다 (3 ≤ m, n ≤ 1000).
• 세 번째 줄부터 차례로 m 개의 줄에는 피자 A의 미리 잘라진 피자조각의 크기를 나타내는 정수가 주어진다.
• 그 다음 n 개의 줄에는 차례로 피자B의 미리 잘라진 피자조각의 크기를 나타내는 정수가 주어진다.
• 각 종류의 피자조각의 크기는 시계방향으로 차례로 주어지며, 각 피자 조각의 크기는 1000 이하의 자연수이다.

출력

• 첫째 줄에는 피자를 판매하는 방법의 가지 수를 나타내는 정수를 출력한다.
• 피자를 판매하는 방법이 없는 경우에는 숫자 0을 출력한다.

# 접근 방법

• 한 종류의 피자에서 2조각 이상을 판매할 때는 반드시 연속된 조각들을 잘라서 판매해야 되는 것이기에 누적합을 떠올렸다.
• 이 때, 하나의 조각에서 나올 수 있는 모든 경우의 수를 저장해주어야 한다.
  • 1 2 3 4 5 의 조각이 있다면 원형으로 되어있기에 1 2 3 5 도 연속된 조각으로 판단할 수 있기 때문이다.
• prefix 배열을 A와 B에 대해 만들어주는데 N을 넘어가면 안되기에 크기는 N+1로 만들어 준다.
• 반복문 2개를 이용하여 모든 경우의 수를 구해주며 예제의 경우
  • 2 2 1 7 2, 2 1 7 2 2, 1 7 2 2 2 와 같이 제일 앞을 돌려주며 경우의 수를 구해주고
  • N의 크기를 넘어선다면 종료, 아니라면 만들어준 prefix 리스트에 경우의 수를 더해준다.
• 출력은
  • A에서 나올 수 있는 N의 개수 + B에서 나올 수 있는 N의 개수 + prefix_A[i] * B에서 [N - i] 곱 을 해주면 된다.
    • 7의 경우 A에서 2, B에서 5를 만드는 것과 같이 N을 만들 수 있고 경우의 수이므로 곱하기로 더해준다.

import sys  
sys.stdin = open('input.txt')  
input = sys.stdin.readline  

def prefix(arr, leng, kind):  
    # 각 피자에 대한 전체 경우의 수 구해주기  
    for i in range(leng):  
        val = 0  
        # 범위의 경우 피자는 원형이므로 제일 앞과 제일 뒤의 조각도 연결되어있다  
        # 따라서 2 2 1 7 2 -> 2 1 7 2 2 와 같이 돌려가며 구해준다.        # 따라서 마지막에 피자 전체크기를 1로 초기화해주는 작업 필요        
        for now in arr[i:] + arr[:i]:  
            val += now  
            # 구해야되는 크기 넘어가면 종료  
            if val > N:  
                break  
            if kind == 0:  
                prefix_A[val] += 1  
            else:  
                prefix_B[val] += 1  
    # 피자 전체 크기는 1번만 나오도록 해주어야 한다.  
    if sum(A) <= N:  
        prefix_A[sum(A)] = 1  
    if sum(B) <= N:  
        prefix_B[sum(B)] = 1  

N = int(input())  
a, b = map(int, input().split())  
A, B = [int(input()) for _ in range(a)], [int(input()) for _ in range(b)]  
prefix_A, prefix_B = [0] * (N+1), [0] * (N+1)  
prefix(A, len(A), 0)  
prefix(B, len(B), 1)  

# 출력은 A와 B에서 각각 독립적으로 N을 만드는 경우의 수를 더해주고  
# 7의 경우 A에서 크기 3, B에서 크기 4를 만드는 경우의 수를 곱해주어야 한다.  
print(prefix_A[N] + prefix_B[N] + sum([prefix_A[i] * prefix_B[N-i] for i in range(N)]))