[Python] 수들의 합2 (백준 2003번 파이썬)

문제

N개의 수로 된 수열 A[1], A[2], …, A[N] 이 있다. 이 수열의 i번째 수부터 j번째 수까지의 합 A[i] + A[i+1] + … + A[j-1] + A[j]가 M이 되는 경우의 수를 구하는 프로그램을 작성하시오.

입력

첫째 줄에 N(1 ≤ N ≤ 10,000), M(1 ≤ M ≤ 300,000,000)이 주어진다. 다음 줄에는 A[1], A[2], …, A[N]이 공백으로 분리되어 주어진다. 각각의 A[x]는 30,000을 넘지 않는 자연수이다.

출력

첫째 줄에 경우의 수를 출력한다.

예제 입력 1 복사

4 2
1 1 1 1

예제 출력 1 복사

3

예제 입력 2 복사

10 5
1 2 3 4 2 5 3 1 1 2

예제 출력 2 복사

3

문제 풀이

시간제한이 0.5초에 수열의 합을 체크하는 문제라서 완전 탐색으로는 해결할 수 없을 것 같다고 생각했습니다.

초기에 투 포인터만을 이용해서 접근했는데, 이후 누적합 방식이 좀 더 좋다는 생각에 접근 방법을 바꿨습니다.

 

포인터의 오른쪽 끝이 배열의 끝에 닿았을 때, 왼쪽  포인터는 아직 중간에 남아있을 수 있으므로 마지막 남은 부분에 대해서도 체크해야합니다. 정렬된 배열이었다면 종료해도 됐겠지만, 정렬된 배열이 아니기 때문에 전부 체크해봐야 합니다.

 

코드

# 2003 수들의 합 2

def get_sum(n, m):

  s, e = 0, 0
  ans = 0
  point_sum = 0

  while e < n:

    # 누적합이 적으면 오른쪽을 늘림
    if point_sum < m:
      point_sum += arr[e]
      e += 1

    # 누적합이 m과 같으면 정답 처리 후 
    elif point_sum == m:
      ans += 1
      point_sum += arr[e]
      e += 1
    # 누적합이 m보다 크면 왼쪽걸 빼야 함 
    else:
      point_sum -= arr[s]
      s += 1
      
  # 마지막 구간 처리해야 함 : s는 중간인데 e가 끝에 닿은 경우,
  # s만 오른쪽으로 움직이면서 마지막도 체크해야 함
  while s < n:
    if point_sum == m:
      ans += 1
    point_sum -= arr[s]
    s += 1

  return ans

N, M = map(int, input().split())
arr = list(map(int, input().split()))
print(get_sum(N, M))