[Python] 수들의 합 : two pointer

수들의 합 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 8 3 1 2 1 3 1 1 1 2

 

▣ 출력예제

1 5

 

 

실패한 풀이

단순 for문과 sum을 사용해서 문제를 풀었으나 시간초과로 실패했다.

sum은 O(n)의 시간복잡도를 가지므로 상당히 낭비이다.

해당 문제처럼 i~j번째의 합이 조건에 맞는지를 검사하는 최적의 방법은 two pointer이다.

합의 결과에 따라 i와 j를 다음 배열 index로 이동시키는 방법이다.

 

해결 풀이

import sys

sys.stdin=open("input.txt", "rt")

n, m = map(int, input().split())
l = list(map(int, input().split()))

count = 0

#index 2개
lt = 0
rt = 1

tot = l[0] #sum(l[lt:rt])


while(True):
    #현재 연산이 목표 값보다 작은 경우
    #오른쪽 포인터를 앞으로 이동시킨다.
    if tot < m:
        #오른족 포인터가 더 이상 갈 수 없다면 종료한다
        if rt == n:
            break
        tot += l[rt]
        rt += 1
    #현재 연산이 목표 값과 같은 경우
    #카운팅하고 왼쪽 포인터를 앞으로 이동시킨다.
    elif tot == m:
        count += 1
        tot -= l[lt]
        lt += 1
    #현재 연산이 목표 값보다 커진 경우
    #왼쪽 포인터를 앞으로 이동시킨다.
    elif tot > m:
        tot -= l[lt]
        lt += 1
    
        
print(count)

즉, twopointer는 연속된 배열에서 특정 구간의 합이 목표 값과 일치하는지를 확인하는 기법이다.