문제

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

www.acmicpc.net/problem/2003

  • 입력

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

  • 출력

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

과정

처음엔 그냥 이중 for loop 돌려서 풀었다가 시간초과가 떴다.

import sys

n, m = map(int, sys.stdin.readline().split())
num = list(map(int, sys.stdin.readline().split()))
res = 0

for i in range(n):
    for j in range(i, n):
        if m == sum(num[i:j+1]):
            res += 1

print(res)

시간을 줄일 방법을 생각하다가 dp를 이용해서 풀어보았다.

dp[i] = 인덱스 0 ~ i 까지의 숫자들의 합으로 놓고 dp[j] - dp[i] 값을 통해 부분합을 구했다. 하지만 이것도 시간초과.

import sys

n, m = map(int, sys.stdin.readline().split())
num = [0]+list(map(int, sys.stdin.readline().split()))
dp = [0]*(n+1)
res = 0

for i in range(n+1):
    dp[i] = sum(num[:i+1])

for i in range(1, n+1):
    for j in range(i, n+1):
        tmp = dp[j] - dp[i-1]
        if m == tmp:
            res += 1
        elif m < tmp:
            break

print(res)

최종 코드

결국 투포인터를 사용했는데 투포인터가 헷갈리는 것이다ㅜㅜ. 기본적인 투포인터가 헷갈려서 슬펐지만 어쨌건 풀긴 풀었다..

근데 다른 코드들을 찾다보니 합이 같은 경우에 포인터 j만 +1 해주는 경우가 많던데 숫자가 모두 자연수라서 i와 j 모두 +1 해주는 게 맞는 것 같다. 뭐 결과는 같겠지만,,

import sys

n, m = map(int, sys.stdin.readline().split())
# index error를 막기 위해 마지막에 0을 추가
num = list(map(int, sys.stdin.readline().split())) + [0]
res = 0
sum = num[0]
i, j = 0, 0

while i < n and j < n:
    # 합이 같은 경우 i와 j를 1씩 증가
    if sum == m:
        res += 1
        sum -= num[i]
        i += 1
        j += 1
        sum += num[j]
    # 합이 작은 경우 j + 1
    elif sum < m:
        j += 1
        sum += num[j]
    # 합이 큰 경우 i + 1
    else:
        sum -= num[i]
        i += 1
        # i가 j를 넘어서는 경우 j도 1 증가
        if i - 1 == j:
            j += 1
            sum += num[j]

print(res)

참고 사이트