[파이썬] 백준 2632 : 피자판매 : (골드2)

[파이썬] 백준 2632 : 피자판매 : (골드2)

• https://www.acmicpc.net/problem/2632
  

---

문제

고객이 두 종류의 피자 A와 B를 취급하는 피자가게에서 피자를 주문하고자 한다.

고객이 원하는 피자의 크기를 이야기하면, 피자가게에서는 한 종류의 피자를 2 조각 이상 판매할 때는 반드시 연속된 조각들을 잘라서 판매한다. 이때 판매한 피자조각의 크기 합이 주문한 크기가 되어야 한다. 판매한 피자조각은 모두 A종류이거나, 모두 B종류이거나, 또는 A와 B 종류가 혼합될 수 있다.

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

입력

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

출력

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

---

풀이

0. 방향성 생각

• 각 피자가 최대 1000조각으로 나누어진다.
• 연속된 피자 조각으로 만들어지는 피자 size들을 모두 저장한다.
• 피자 사이즈가 1개 ~ 1000개일때 환형으로 모두 체크해야하므로 $1000^2$이다.
• 피자 조합들을 저장한 자료구조의 크기가 10만이므로, A에서 B피자로 매칭 시킬 때 전부 조사가 불가능하다.
• 정렬 후 이분탐색으로 $NlogN$으로 풀이.

0. 누적합

1. 피자 입력

from collections import defaultdict as dd
import sys
input = lambda : sys.stdin.readline().rstrip()

size = int(input())
m,n = map(int,input().split())

def pizza(p):
    T = [int(input())]
    for _ in range(p-1):
        T.append(T[-1]+int(input()))
    return T
A = pizza(m)
B = pizza(n)

• 피자 1개의 크기는 별로 중요하지 않다. 바로 누적합으로 저장 (1개 사이즈는 인접 인덱스에서 빼면 된다)
  

2. 연속된 피자 경우의 수 저장

def pizza_comb(T):
    comb = dd(int)
    total = T[-1]
    tl = len(T)
    for i in range(tl): # 피벗
        for j in range(1,tl): # 조각 연속
            if i+j >= tl:
                comb[T[(i+j)%tl] + total - T[i]] += 1
            else:
                comb[T[i+j]-T[i]] += 1
    comb[total] = 1
    comb = list(comb.items())
    comb.sort()
    return comb

Acomb = pizza_comb(A)
Acomb = [(0,1)] + Acomb
Bcomb = pizza_comb(B)

• 어떤 피자 조각 i를 피벗으로 만들고, 크기가 1,2,...,피자 조각 수에 따라서 총 조각 사이즈를 저장한다.
• 이분 탐색을 해야하므로 리스트로 만들고 정렬해서 리턴한다.
• A 피자를 선택하지 않고, B 피자를 만드는 경우가 있으므로 (0,1)이라는 튜플도 추가한다.
• comb[연속된 피자의 총 사이즈] = 경우의 수
  

3. 이분탐색

bl = len(Bcomb)
answer = 0
for a,count in Acomb: # 피자 사이즈 a와 경우의 수 count
    goal = size-a
    if goal == 0: # A만으로도 해결이 가능한 경우
        answer += count
    elif goal > 0: # A를 고르고 B 피자에서 가져와야 하는 경우
        l,r = 0,bl -1
        find = False
        while l <= r:
            m = l + (r-l)//2
            if Bcomb[m][0] == goal:
                find = True
                break
            elif Bcomb[m][0] < goal:
                l = m+1
            else:
                r = m-1
        if find: # 찾은 경우에는 각 경우의 수의 곱
            answer += count * Bcomb[m][1]
print(answer)

• 피자 사이즈 a와 경우의 수 count에 대해서 B 피자를 탐색한다.
• 피자 A가 size를 만족시키면 바로 다음 피자
• 만족시키지 못하면 B에서 이분탐색을 실행한다
  

---

전체코드

from collections import defaultdict as dd
import sys
input = lambda : sys.stdin.readline().rstrip()

size = int(input())
m,n = map(int,input().split())

def pizza(p):
    T = [int(input())]
    for _ in range(p-1):
        T.append(T[-1]+int(input()))
    return T
A = pizza(m)
B = pizza(n)

def pizza_comb(T):
    comb = dd(int)
    total = T[-1]
    tl = len(T)
    for i in range(tl): # 피벗
        for j in range(1,tl): # 조각 연속
            if i+j >= tl:
                comb[T[(i+j)%tl] + total - T[i]] += 1
            else:
                comb[T[i+j]-T[i]] += 1
    comb[total] = 1
    comb = list(comb.items())
    comb.sort()
    return comb

Acomb = pizza_comb(A)
Acomb = [(0,1)] + Acomb
Bcomb = pizza_comb(B)

bl = len(Bcomb)
answer = 0
for a,count in Acomb:
    goal = size-a
    if goal == 0:
        answer += count
    elif goal > 0:
        l,r = 0,bl -1
        find = False
        while l <= r:
            m = l + (r-l)//2
            if Bcomb[m][0] == goal:
                find = True
                break
            elif Bcomb[m][0] < goal:
                l = m+1
            else:
                r = m-1
        if find:
            answer += count * Bcomb[m][1]
print(answer)

코멘트

• 환형이라서 인덱스 %를 해주기.
• 한쪽 끝에서 나와서 다른 쪽 끝에서 연결되는 경우는 % 사용하기