시간복잡도 줄이는 기술1. prefix (2) 부분합이 K

문제

n개의 정수로 이루어진 수열에서 연속된 구간의 합을 구하려합니다.

모든 연속된 구간의 합 중에서 합이 k인 것의 개수를 구하는 프로그램을 작성하세요.

 

아이디어

부분합 수열 (prefix)를 구하고, 구간별로 부분합을 구한다.

구한 값이 K보다 크다면 정답이 없다는 뜻이므로, 탐색을 멈추고

구한 값이 K라면 정답+1 

구한 값이 K보다 작다면 뒷 구간 탐색을 계속 진행하며 더하는 값을 늘려준다. 

 

코드

import sys 
n,k=map(int,sys.stdin.readline().split()) 
arr=[0]+ list(map(int,sys.stdin.readline().split()))

prefix_sum =[0]*(n+1)
for i in range(1,n+1):
    prefix_sum[i]=arr[i]+prefix_sum[i-1]

ans=0
for i in range(n+1):
    for j in range(i+1,n+1):
        #N의 조건이 양수므로 부분합이 K보다 크거나 같으면 뒤는 더 탐색할 필요 X
        if prefix_sum[j]-prefix_sum[i] > k:
            break
        elif  prefix_sum[j]-prefix_sum[i]== k:
            ans+=1
            break
print(ans)