[프로그래머스] 연속된 부분 수열의 합 Lv.2 Python

문제 링크

1. 문제 설명

비내림차순으로 정렬된 수열이 주어질 때, 다음 조건을 만족하는 부분 수열을 찾으려고 합니다.

• 기존 수열에서 임의의 두 인덱스의 원소와 그 사이의 원소를 모두 포함하는 부분 수열이어야 합니다.
• 부분 수열의 합은 k입니다.
  합이 k인 부분 수열이 여러 개인 경우 길이가 짧은 수열을 찾습니다.
• 길이가 짧은 수열이 여러 개인 경우 앞쪽(시작 인덱스가 작은)에 나오는 수열을 찾습니다.

수열을 나타내는 정수 배열 sequence와 부분 수열의 합을 나타내는 정수 k가 매개변수로 주어질 때, 위 조건을 만족하는 부분 수열의 시작 인덱스와 마지막 인덱스를 배열에 담아 return 하는 solution 함수를 완성해주세요. 이때 수열의 인덱스는 0부터 시작합니다.

입출력 예

sequence	k	result
[1, 2, 3, 4, 5]	7	[2, 3]
[1, 1, 1, 2, 3, 4, 5]	5	[6, 6]
[2, 2, 2, 2, 2]	6	[0, 2]

2. 문제 접근

시작 인덱스와 끝 인덱스를 동시에 구해야 하는 문제입니다. 수열이 사실상 오름차순으로 정렬되어 있어 매우 직관성이 좋네요. 끝 부분부터 탐색해도 좋지만 여러 개의 가능한 수열 후보 중에서 길이가 같다면 시작이 가장 빠른 수열을 찾아야 하니 구현이 조금 복잡합니다.

따라서 저는 시작 부분부터 탐색을 시작했습니다. 시작 인덱스에 해당하는 start에서 시작하여 end까지 수열의 부분합이 k보다 크거나 같을 때까지 탐색합니다. 만약에 수열의 부분합이 k보다 크다면 start에서 시작한 수열의 부분합으로 k값을 만들수 없다고 판단하여 다음 start부터 탐색을 다시 시작합니다. 부분합이 k인 다른 부분 수열을 수열의 뒷 부분을 탐색하다가 찾았다면 새로 찾은 부분 수열의 길이가 기존보다 짧을 때 새로운 수열의 start, end가 문제가 요구하는 정답이 됩니다.

또한 k가 부분 수열에서 만들 수 없는 경우도 존재하지 않는다고 문제에서 명시되어 있으므로 별도의 예외 처리도 필요하지 않습니다. 무한 루프에 빠질 걱정 없이 반복문을 돌리실 수 있습니다.

3. 문제 풀이

def solution(sequence, k):
    if sum(sequence) < k: # 풀이가 불가능하면
        return [-1, -1]

    start, end = 0, 0
    cur_sum = 0
    answer = [0, len(sequence)]

    while end < len(sequence):
        cur_sum += sequence[end]

        while start <= end and cur_sum > k: #부분합이 k를 넘는다면
            cur_sum -= sequence[start] 
            start += 1 #다음 start를 탐색

        if cur_sum == k and end - start < answer[1] - answer[0]:
            answer = [start, end]

        end += 1

    return answer