[백준] 수들의 합2(python)

백준, 수들의 합 2

TL;DR

• 투 포인터(two pointer)

문제 분석

1. N개의 자연수 배열에 대해서 i부터 j까지 연속된 수들의 합이 M이 되는 경우의 수를 구하는 프로그램을 작성하라.
2. 1 <= N <= 10,000 1 <= M <= 300,000,000. 배열의 원소는 30,000을 넘지 않는 자연수이다.

 

- 주어진 자연수 배열에서 i번째 인덱스부터 j번째 인덱스까지 연속된 원소들의 합이 M이 되는 경우의 수를 반환하면 된다.

입출력 형태

입출력 예시

 

- 예시 1번은 [1, 1, 1, 1]로 이루어진 길이가 4인 배열이 주어지고, 이 배열에서 합이 2가 되는 연속된 배열 경우의 수를 구하면 된다.

  - 이 경우 0번 인덱스 ~ 1번 인덱스[1, 1], 1번 인덱스 ~ 2번 인덱스[1, 1], 2번 인덱스 ~ 3번 인덱스[1, 1] 총 3개의 경우가 존재한다.

- 예시 2번에서는 1 ~ 2[2, 3], 5[5], 6 ~ 8[3, 1, 1] 총 3개의 경우가 존재한다.

- 연속된 수들의 합이라는 조건을 염두해두어야 한다.

풀이

import sys

n, m = map(int, sys.stdin.readline().split())
count = 0
arr = list(map(int, sys.stdin.readline().split()))

start, end = 0, 1

while end <= len(arr):
    _sum = sum(arr[start:end + 1])
    if _sum == m:
        count += 1
        start += 1
        end += 1
    
    if _sum < m:
        end += 1

    if _sum > m:
        start += 1

print(count)

 

문제의 조건에서 배열의 길이가 최대 10,000이므로 이 문제는 brute force로 해결할 수 있다.

- 두 개의 포인터 start와 end를 사용하여 두 구간 사이에 있는 배열의 원소의 합을 구한다.

  - 구간에서 배열의 합이 m과 일치할 때는 count를 증가시킨다.

  - 구간에서 배열의 합이 m보다 작을 때는 end를 증가시켜 배열의 합을 크게 만든다.

  - 구간에서 배열의 합이 m보다 클 때는 start를 증가시켜 배열의 합을 작게 만든다.

 

이렇게 문제를 해결할 수 있으나 시간 복잡도가 매우 높다. 채점 경우에 따라 다르겠지만 604ms가 나왔다.

개선된 풀이

위 방법은 불필요하게 많은 원소들을 조회한다는 문제가 있다. 어떻게 개선할 수 있을까 ?

import sys

n, m = map(int, sys.stdin.readline().split())
arr = list(map(int, sys.stdin.readline().split()))

_sum, index, count = 0, 0, 0

for item in arr:
    _sum += item
    
    while _sum > m:
        _sum -= arr[index]
        index += 1
        
    count += (_sum == m)

print(count)

 

이 코드는 주어진 배열을 한 번만 순회해도 된다.

- 배열을 순회하며 원소를 방문하고, 방문한 원소를 _sum에 누적시킨다.

- 그러고 나서는 _sum 원소와 주어진 m을 비교한다.

- 만약 _sum이 m보다 크다면 원소를 더 순회해도 절대 m이 될 수 없으므로 시작 인덱스를 제거하고, 시작 인덱스를 다음 인덱스로 둔다.

- 누적 값인 _sum과 m이 같다면 count를 증가시킨다.

 

이 코드는 두 개의 포인터를 사용하는 대신 부분 배열의 시작 인덱스인 index를 사용하여 연속한 배열의 합을 확인한다. 이 경우에는 시간 복잡도가 72ms가 나온다. 이전 코드에 비하여 훨씬 빠르게 실행된다.

투 포인터로 사용할 수 있는 문제에 대해서 무조건 투 포인터로 문제를 해결하는 것이 최고의 풀이가 아님을 확인할 수 있다.