[코드트리 파이썬] 바구니 안의 사탕

문제

 

1차원 그래프에 바구니가 놓여있고 해당 바구니에 사탕이 놓여있다. 이 때 한 지점을 잡아 주어진 범위에 사탕이 가장 많은 경우를 구한다.

---

풀이

N = 100000, K = 2000000 이기 때문에 이중 이상의 루프를 돌면 시간이 초과된다.

따라서 투포인터를 이용하여 해결하였다.

N, K = map(int, input().split())

maxSize = 1000001
arr = [0] * maxSize
basket = set()
for i in range(N):
    a, b = map(int, input().split())
    arr[b] += a
    basket.add(b)

배열의 최대 크기를 정의하고 해당 크기만큼의 배열을 만든다.

그 후 좌표와 사탕의 개수를 입력받아 배열에 더해준다.

이렇게 하는 이유는 같은 위치에 여러 바구니가 있을 수 있기 때문이다.

 

j = 0
k = 0
jSum = 0
kSum = 0
res = 0

for i in range(maxSize):
    while j < i - K:
        jSum += arr[j]
        j += 1
    
    while k <= i + K and k < maxSize:
        kSum += arr[k]
        k += 1
    
    res = max(res, kSum - jSum)

실제 투포인터를 구현하는 부분이다.

0부터 배열의 끝까지 진행하는 포인터와

그 안에 2개의 포인터로 구성된다.

하나는 i-K-1까지의 사탕의 누적개수를 계산하는 포인터와

또 하나는 i + K 까지의 사탕의 누적개수를 계산하는 포인터이다.

두개의 포인터로 누적개수를 구한 후 서로의 차를 구하면 i-K부터 i+K까지의 사탕의 개수가 된다.

 

배열의 모든 부분을 순회하기 때문에 res값을 항상 갱신해주면 가장 큰 값을 구할 수 있다.

 

코드

N, K = map(int, input().split())

maxSize = 1000001
arr = [0] * maxSize
basket = set()
for i in range(N):
    a, b = map(int, input().split())
    arr[b] += a
    basket.add(b)

j = 0
k = 0
jSum = 0
kSum = 0
res = 0

for i in range(maxSize):
    while j < i - K:
        jSum += arr[j]
        j += 1
    
    while k <= i + K and k < maxSize:
        kSum += arr[k]
        k += 1
    
    res = max(res, kSum - jSum)
    
print(res)