문제

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

그림1

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

그림2

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

www.acmicpc.net/problem/2632

  • 입력

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

  • 출력

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

과정

앞서 다른 문제를 이분탐색으로 문제를 풀었기때문에 이 문제도 같은 방식으로 접근했다. 하지만 Counter를 사용해서인지, 가능한 모든 합을 구하는 데가 쓸데없이 복잡해서 시간을 많이 뺏긴 건지 pypy3에서만 통과하고 python3에서는 통과하지 못했다.

import sys
from collections import Counter

s = int(sys.stdin.readline())
m, n = map(int, sys.stdin.readline().split())
a = [int(sys.stdin.readline()) for _ in range(m)]
b = [int(sys.stdin.readline()) for _ in range(n)]

pizzaA = [0, sum(a)]
pizzaB = [0, sum(b)]
# pizza A에서 가능한 합을 모두 구함
for i in range(1, m):
    for j in range(m):
        tmp = 0
        l = j
        for _ in range(i):
            tmp += a[l]
            l += 1
            # index를 넘어가면 처음으로 돌아옴
            if l >= m:
                l -= m
        pizzaA.append(tmp)
# pizza B에서 가능한 합을 모두 구함
for i in range(1, n):
    for j in range(n):
        tmp = 0
        l = j
        for _ in range(i):
            tmp += b[l]
            l += 1
            if l >= n:
                l -= n
        pizzaB.append(tmp)

# Counter 모듈을 활용하여 경우의 수를 구함
counter = Counter(pizzaA)
res = 0
for x in pizzaB:
    res += counter[s-x]

print(res)

최종 코드

다른 코드를 참고해보니 우선 두 피자에서 가능한 합을 구하는 부분이 너무 복잡해서 간결하게 고쳤다. 마치 사전처럼 리스트의 index가 합이 되고, 리스트의 값에 합이 나오는 경우의 수를 넣어서 계산에 용이하도록 했다. 그리고 합이 s를 초과하는 경우에는 합을 저장하지 않고 중단했다.

그리고 index를 모두 돌면서 합이 s가 되는 경우를 결과 값에 더해주어 답을 찾을 수 있었다.

import sys

s = int(sys.stdin.readline())
m, n = map(int, sys.stdin.readline().split())
a = [int(sys.stdin.readline()) for _ in range(m)]
b = [int(sys.stdin.readline()) for _ in range(n)]

pizzaA = [1] + [0] * s
pizzaB = [1] + [0] * s
for i in range(m):
    tmp = 0
    for j in range(m):
        # index를 (i+j) % m으로 접근하여 index error를 방지
        tmp += a[(i+j) % m]
        # tmp > s인 경우는 필요가 없으므로 break
        if tmp > s:
            break
        else:
            pizzaA[tmp] += 1
for i in range(n):
    tmp = 0
    for j in range(n):
        tmp += b[(i+j) % n]
        if tmp > s:
            break
        else:
            pizzaB[tmp] += 1
# count 중복을 방지하기 위해 sum(a) <= s인 경우 피자 전체의 합은 1로 고정
if sum(a) <= s:
    pizzaA[sum(a)] = 1
if sum(b) <= s:
    pizzaB[sum(b)] = 1

res = 0
for i in range(s + 1):
    res += (pizzaA[i] * pizzaB[s - i])

print(res)

참고 사이트