프로그래머스 | 연속된 부분 수열의 합 [파이썬 python]

프로그래머스

 

시간 초과^^

def solution(sequence, k):
    # 합 배열
    s = [0] * (len(sequence)+1) 
    for n in range(1, (len(sequence)+1)):       
        s[n] = s[n-1] + sequence[n-1]
    
    # 구간 합
    answer = []
    min_len = 1000001
    for i in range(len(sequence)+1):
        for j in range(i, len(sequence)+1):
            if s[j] - s[i-1] == k and j-i < min_len:
                if answer:
                    answer.pop()
                answer.append([j-i, [i-1,j-1]])   
                min_len = j-i
                
    answer.sort(key=lambda x: [x[0], x[1]])
    return answer[0][1]

문제 이름을 보자마자 아! 구간 합 문제구나! 하고 호기롭게 코드를 써 내려갔지만 결과는 처참하게도 시간 초과..^^

그도 그럴 것이 5 ≤ sequence의 길이 ≤ 1,000,000 이기 때문에 for문에서 엄청난 시간이 걸릴 것이다.

 

정답 코드

def solution(sequence, k):
    start, end = 0, 0 # 시작 인덱스, 마지막 인덱스
    prefix_sum = sequence[0] # 부분 수열의 합
    min_len = 1000001 # 최소 길이
    
    while start <= end < len(sequence):
        # 부분 수열의 합이 k인 경우
        if prefix_sum == k:
            # min_len보다 길이가 짧은 수열이면
            if end-start+1 < min_len:
                min_len = end-start+1 # 최소 길이 갱신
                answer = [start, end] # answer 갱신
            
            # 다음 작업을 위해 부분 수열의 합에서 start 값을 빼주고, 시작 인덱스 +1
            prefix_sum -= sequence[start]
            start += 1
            
        # 부분 수열의 합이 k보다 작은 경우
        elif prefix_sum < k:
            end += 1 # 마지막 인덱스 +1
            # 마지막 인덱스가 sequence 범위 안에 있으면 부분 수열의 합에 end 누적
            if end < len(sequence):
                prefix_sum += sequence[end]
                
        # 부분 수열의 합이 k보다 큰 경우
        elif prefix_sum > k:
            # 부분 수열의 합에서 start를 빼주고 start +1
            prefix_sum -= sequence[start]
            start += 1
            
    return answer

아무리 생각해도 어떻게 시간 초과를 해결해야 할지 몰라서 힌트를 찾아다닌 결과 투포인터로 풀어야 한다는 것을 알게 되었다.

 

start, end 인덱스를 변수로 설정해 부분 sequence[start:end]와 k를 비교해 포인터 위치를 +1 또는 -1을 해주어야 하는 것이다.

sequence의 처음부터 끝까지 탐색할 것이므로 while문을 통해 start와 end가 len(sequence)를 벗어나면 반복을 중단하도록 했다.

 

그다음으로 구간 합과 k를 비교한다.

부분 수열의 합이 k인 경우, 즉 조건에 부합하는 경우엔 min_len과 비교하여 더 짧은 길이인 경우에만 answer로 변경해 준다.

다음 작업(=[0:2]를 탐색했다면, [1:2]를 탐색)을 위해 start 위치를 +1 한다.

# 부분 수열의 합이 k인 경우
if prefix_sum == k:
    # min_len보다 길이가 짧은 수열이면
    if end-start+1 < min_len:
        min_len = end-start+1 # 최소 길이 갱신
        answer = [start, end] # answer 갱신

    # 다음 작업을 위해 부분 수열의 합에서 start 값을 빼주고, 시작 인덱스 +1
    prefix_sum -= sequence[start]
    start += 1

부분 수열의 합이 k보다 작은 경우에는 부분 수열에 큰 값을 더 추가해 k에 가까워질 수 있도록 한다.

그러므로 end+1을 진행하고, end의 범위를 확인한다.

len(sequence) 안에 있으면 sequence에 인덱스로 접근할 수 있으므로 해당 값을 부분 수열 합에 누적시켜 준다.

범위 안에 존재하지 않으면 그 값은 범위보다 큰 것이다.

예를 들어, len(sequence)=5일 때 end=6이면 [start:end] 인덱싱은 어차피 len(sequence)까지만 된다. 그러므로 이미 누적합은 구해져 있을 것이므로 더 이상 값을 누적시킬 필요가 없다.

# 부분 수열의 합이 k보다 작은 경우
elif prefix_sum < k:
    end += 1 # 마지막 인덱스 +1
    # 마지막 인덱스가 sequence 범위 안에 있으면 부분 수열의 합에 end 누적
    if end < len(sequence):
        prefix_sum += sequence[end]

마지막으로, 부분 수열의 합이 k보다 큰 경우에는 start+1을 통해 범위를 좁혀나간다.

# 부분 수열의 합이 k보다 큰 경우
elif prefix_sum > k:
    # 부분 수열의 합에서 start를 빼주고 start +1
    prefix_sum -= sequence[start]
    start += 1

이렇게 코드를 제출하면 시간초과가 발생하던 테스트 케이스들에서도 무리 없이 통과된다.