Subarray Sum Equals K

+10

정수 배열 nums와 정수 k가 주어집니다. 원소의 합이 정확히 k인 부분 배열의 개수를 세세요. 부분 배열은 하나 이상의 인접한 원소로 이루어진 연속 구간입니다. 값이 같더라도 시작 위치나 끝 위치가 다르면 두 부분 배열은 각각 별도로 셉니다. 값은 음수이거나 0일 수 있습니다.

함수

subarraySum(nums: integer-array, k: integer) → integer

numsinteger-array: 음수와 0을 포함할 수 있는 정수 배열
kinteger: 부분 배열로 간주되기 위해 도달해야 하는 합
반환값integer: 요소의 합이 k가 되는 부분 배열의 개수

제약 조건

1 ≤ nums.length ≤ 2 × 10⁴
-1000 ≤ nums[i] ≤ 1000
-10⁷ ≤ k ≤ 10⁷
이렇게 긴 배열의 부분 배열은 최대 200,010,000개이므로, 답은 32비트 부호 있는 정수에 들어갑니다.

예제

입력: nums = [3, 4, -7, 1, 3, 3, 1, -4]k = 7
출력: 4
설명: 네 개의 연속 구간의 합은 7입니다: [3, 4], [1, 3, 3], [3, 3, 1] 및 [3, 4, -7, 1, 3, 3]. 마지막 구간에서는 -7이 3과 4를 상쇄하고, 나중에 합이 다시 7이 되므로 합이 k를 초과한 뒤에도 구간이 일치할 수 있습니다.

입력: nums = [1, -1, 0]k = 0
출력: 3
설명: 합이 0인 부분 배열은 세 개입니다: [1, -1], [0], 그리고 전체 배열인 [1, -1, 0]입니다. [-1, 0] 구간의 합은 -1이므로 포함되지 않습니다.

입력: nums = [2, 2, 2]k = 4
출력: 2
설명: 인덱스 0과 1에 있는 [2, 2] 연속 구간과 인덱스 1과 2에 있는 [2, 2] 연속 구간은 같은 값을 포함하지만 서로 다른 위치에 있으므로 둘 다 셉니다. 배열 전체의 합은 6입니다.

제출 시 숨은 테스트 +17개

후속 질문

합계가 k가 되는 가장 긴 부분 배열의 길이를 O(n) 시간 안에 반환하도록 해법을 어떻게 바꾸시겠어요?

힌트

하나씩 열어 보세요. 열 때마다 조금 더 알려 줍니다.

3개 중 0개 열림

풀이

n개의 숫자로 이루어진 배열에는 n(n+1)/2개의 부분 배열이 있으며, n = 2 × 10^4일 때 약 2 × 10^8개이므로 각 부분 배열의 합을 구하는 것은 너무 느립니다. 음수 값 때문에 슬라이딩 윈도우도 사용할 수 없습니다. 윈도우의 합은 감소했다가 다시 증가할 수 있으므로, 언제 윈도우를 줄여야 하는지 알려 주는 규칙이 없습니다. 문제를 해결하는 핵심 아이디어는 모든 부분 배열의 합을 두 접두사 합의 차이로 나타내는 것입니다. 현재 원소에서 끝나고 합이 k인 부분 배열의 개수를 세는 것은 현재 접두사 합에서 k를 뺀 값과 같은 이전 접두사 합의 개수를 세는 것을 의미하며, 해시 맵을 사용하면 한 번의 순회로 이를 계산할 수 있습니다.

실행 중인 총합으로 모든 것을 시작하세요

정답이지만 가장 큰 테스트에서는 끝나지 않음

핵심 아이디어

모든 부분 배열에는 첫 번째 인덱스 start와 마지막 인덱스 end가 있습니다. 모든 쌍을 방문해 합을 확인하면 각 부분 배열을 정확히 한 번씩 확인하게 되므로 개수가 정확합니다.

각 부분 배열의 합을 구하기 위해 세 번째 반복문을 사용할 필요는 없습니다. start를 고정한 다음 end를 한 번에 한 단계씩 오른쪽으로 옮기면서 nums[end]를 누적 total에 더하세요. total에는 항상 start부터 end까지 원소들의 합이 저장되므로, 각 부분 배열마다 덧셈 한 번과 비교 한 번이면 됩니다.

합이 k에 도달하거나 이를 넘어섰다고 해서 멈추지 마세요. 나중에 음수가 나오면 합이 다시 내려갈 수 있습니다. 첫 번째 예시에서 인덱스 0부터 시작하는 합은 3, 7, 0, 1, 4, 7이 되므로, 인덱스 5에서 두 번째로 일치합니다.

비용은 쌍의 개수입니다. n = 2 × 10^4일 때 쌍은 약 2 × 10^8개이며, C에서는 괜찮지만 Python, Ruby 또는 R에서는 너무 느립니다.

알고리즘

count를 0으로 설정합니다.
0부터 n-1까지 각 start에 대해 total을 0으로 설정합니다.
start부터 n-1까지 각 end에 대해 nums[end]를 total에 더합니다.
total이 k와 같으면 count에 1을 더하고, 어느 쪽이든 계속 진행합니다.
count를 반환합니다.

def subarraySum(nums, k):
    n = len(nums)
    count = 0
    for start in range(n):
        total = 0
        # Grow the subarray that begins at start, one element at a time
        for end in range(start, n):
            total += nums[end]
            if total == k:
                count += 1
    return count

개수 맵을 사용한 누적 합

핵심 아이디어

prefix[j]를 처음 j개 요소의 합이라고 하며, 빈 접두 부분 배열의 경우 prefix[0] = 0입니다. 인덱스 i부터 인덱스 j-1까지의 부분 배열의 합은 prefix[j] - prefix[i]입니다. 따라서 현재 요소에서 끝나는 부분 배열의 합이 k가 되는 것은 이전 접두 합이 현재 접두 합에서 k를 뺀 값과 같을 때입니다. 이처럼 일치하는 이전 접두 합은 각각 일치하는 부분 배열이 시작하는 위치를 나타냅니다.

배열을 한 번 순회합니다. 누적 접두 합과 각 접두 합이 나타난 횟수를 저장하는 해시 맵 seen을 유지합니다. 각 요소에서 먼저 seen[prefix - k]를 개수에 더한 다음 현재 접두 합을 기록합니다. 기록하기 전에 조회하면 부분 배열이 비어 있는 경우를 막을 수 있습니다. k = 0일 때 먼저 기록하면 현재 접두 합과 그 자체가 일치하게 됩니다.

k = 7인 첫 번째 예를 살펴보겠습니다. 접두 합은 0, 3, 7, 0, 1, 4, 7, 8, 4입니다. 인덱스 1 이후 접두 합이 7에 도달할 때 맵에는 0이 하나 있으므로 [3, 4]가 나옵니다. 인덱스 5 이후 접두 합이 다시 7에 도달할 때 맵에는 빈 접두 부분 배열과 -7 이후의 접두 부분 배열에 해당하는 0이 두 개 있으며, 이 두 접두 합은 각각 [3, 4, -7, 1, 3, 3]과 [1, 3, 3]을 한 번에 찾아냅니다. 인덱스 6 이후 8에 도달할 때 맵에는 1이 하나 있으므로 [3, 3, 1]이 나옵니다. 따라서 총 4개입니다.

맵에 0이 한 번 나타났다고 초기화하면 인덱스 0에서 시작하는 부분 배열을 셀 수 있습니다. 집합이 아니라 횟수를 저장하는 맵이 중요한 이유는 같은 접두 합이 반복될 수 있고, 각 항목이 서로 다른 부분 배열을 시작하기 때문입니다. 각 요소마다 조회 한 번과 갱신 한 번이 필요하므로 시간 복잡도는 O(n)이고, 맵에는 최대 n+1개의 키가 저장됩니다.

알고리즘

seen 맵을 만들고 seen[0] = 1로 설정한 다음, prefix와 count를 0으로 설정합니다.
각 요소를 prefix에 더합니다.
seen[prefix - k]를 count에 더합니다. 키가 없으면 0으로 읽습니다.
seen[prefix]에 1을 더합니다.
count를 반환합니다.

def subarraySum(nums, k):
    # seen[p] = how many prefixes so far add up to p; the empty prefix adds up to 0
    seen = {0: 1}
    prefix = 0
    count = 0
    for num in nums:
        prefix += num
        # Each earlier prefix equal to prefix - k starts a subarray that ends here and sums to k
        count += seen.get(prefix - k, 0)
        seen[prefix] = seen.get(prefix, 0) + 1
    return count

함정과 경계 사례

대부분의 오답은 모든 값이 양수라고 가정하거나, 두 맵 연산의 순서를 잘못 처리해서 발생합니다.

합이 k를 넘으면 축소하는 슬라이딩 윈도우는 음수 값이 있을 때 실패합니다. 첫 번째 예시에서 4가 아니라 2를 반환합니다. 합이 인덱스 6에서 7을 넘을 때까지 윈도우의 왼쪽 끝을 인덱스 0에 그대로 두므로, [1, 3, 3]이나 [3, 3, 1]을 확인하지 않습니다.
seen[0] = 1을 빠뜨리면 인덱스 0에서 시작하는 모든 부분 배열을 놓칩니다. nums = [5]와 k = 5의 경우 1이 아니라 0을 반환합니다.
조회 전에 현재 접두사 합을 기록하면 k가 0일 때 빈 부분 배열도 세게 됩니다. [1, -1, 0]의 경우 3이 아니라 6을 반환합니다.
횟수 맵 대신 접두사 합 집합을 사용하면 반복되는 값을 제대로 세지 못합니다. [0, 0, 0]과 k = 0의 경우 정답은 6입니다. 같은 접두사 합이 앞서 나온 각각의 경우가 서로 다른 부분 배열을 시작하기 때문입니다.
브루트 포스에서는 합이 k를 넘었을 때 내부 루프를 중단하는 것도 슬라이딩 윈도우와 같은 이유로 잘못된 방법입니다.

자주 묻는 질문4

Subarray Sum Equals K의 시간 복잡도는 무엇인가요?

접두사 합과 해시 맵을 사용하는 풀이는 O(n) 시간과 O(n)의 추가 공간이 필요합니다. 한 번 순회하면서 각 요소마다 조회 한 번과 갱신 한 번을 수행합니다. 누적 합을 사용해 모든 부분 배열을 확인하면 O(n²) 시간이 걸리고, 각 부분 배열의 합을 처음부터 다시 계산하면 O(n³) 시간이 걸립니다.

Subarray Sum Equals K에서 슬라이딩 윈도우가 작동하지 않는 이유는 무엇인가요?

슬라이딩 윈도우는 윈도우가 커지면 합이 증가하고 작아지면 감소한다는 원리에 의존하며, 이는 모든 값이 양수일 때만 성립합니다. 음수 값이 있으면 합이 이미 너무 큰 윈도우도 더 커진 후에 조건을 만족할 수 있으므로, 왼쪽 경계를 언제 이동해야 하는지 알려 주는 규칙이 없습니다. 모든 값이 양수라면 슬라이딩 윈도우로 O(n) 시간과 O(1) 공간에 문제를 해결할 수 있습니다.

해시 맵은 왜 0을 1에 매핑한 상태로 시작하나요?

그 항목은 첫 번째 요소 앞의 빈 접두사를 나타내며, 그 합은 0입니다. 인덱스 0에서 시작하는 부분 배열의 합은 현재 접두사 합에서 이 빈 접두사를 뺀 값이므로, 이 항목이 없으면 해당 부분 배열은 절대 세어지지 않습니다. nums = [5]와 k = 5의 경우, 5 - 5 = 0을 조회하면 해당 항목을 찾아 1을 반환합니다.

부분 배열의 합이 K와 같은 문제를 추가 공간 O(1)로 해결할 수 있을까요?

한 번만 순회하는 방법으로는 불가능합니다. 어떤 원소에서 끝나는 일치 항목의 수를 세려면 그 원소보다 앞에 어떤 접두 합이 있었는지 알아야 하며, 서로 다른 접두 합이 최대 n+1개일 수 있습니다. 맵이 없으면 O(n²)의 누적 합 방식으로 돌아가게 됩니다. 모든 값이 양수일 때는 슬라이딩 윈도우를 사용해 O(n) 시간과 O(1) 공간으로 부분 배열을 셀 수 있습니다.

Python

def subarraySum(nums, k):
    # 여기에 코드를 작성하세요

테스트 케이스

케이스 1

케이스 2

케이스 3

입력

nums = [3, 4, -7, 1, 3, 3, 1, -4]
k = 7

기대값