문제

N개의 정수가 담긴 배열 A와 정수 S가 주어집니다. 이때 철수는 평균이 S가 되는 A에서의 부분 수열의 개수를 구하고 싶다고 합니다.
만약 결과가 10억보다 크다면 10억을 리턴해주세요.

예시

• A = [2, 1, 3]이고 S = 2일때, 평균이 2가 되는 A의 부분수열은 [2], [1, 3], [2, 1, 3]으로 총 3개 입니다.
• A = [0, 4, 3, -1]이고 S = 2일때, 평균이 2가 되는 A의 부분수열은 [0, 4], [4, 3, -1]으로 총 2개 입니다.
• A = [2, 1, 4]이고 S = 3일 때, 평균이 3이 되는 A의 부분수열은 없으므로 0을 리턴합니다.

제한 사항

• N은 1 이상 10만 이하입니다.
• S는 -10억 이상 10억 이하의 정수입니다.
• A의 원소는 -10억에서 10억까지의 정수 값을 가질 수 있습니다.

풀이

아이디어

문제에 주어진 평균의 의미를 다시 생각해보면 아래 필기와 같다.

따라서, B라는 배열을 다음과 같이 각 원소에 평균을 뺀 값들로 채워준다.

B = [a - S for a in A]

그 다음에는 어떻게 하는가..?
A = [2, 1, 3]과 S = 2를 가지고 예를 들어보겠다. 이 배열을 가지고 B를 생성했을 때 B는 [0, -1, 1]이 된다.
이때 우리는 A의 부분수열 k개에 대해서 해당 원소들의 합이 0이 되는 순간들을 찾고 싶다.

그러면 다음과 같이 누적합을 활용한 아이디어를 이용하면 된다.
A = [8, -3, 4, 5, -9, -3, 5, -2, 1] 일때, A의 부분수열 중에서 합이 5가 되는 경우를 찾고 싶다고 하자. 그러면 다음과 같이 누적합이 구해질 것이다.

이때, 누적합이 5가 되는 부분들을 보자. 누적합이 5가 되는 부분들이 총 3번 나오는데, 그러면 어떤 5와 다음번 5 사이 원소들의 합이 0이라는 뜻이 된다.


따라서, 누적합이 5가 되는 것들 중에 2개를 골라 출발점, 도착점으로 설정했을 때, 출발점+1부터 도착점까지의 A의 원소를 모두 더하면 0이 되는 것을 확인할 수 있다.

그럼 차근차근 순서대로 했을 때 어떤 규칙을 발견할 수 있는지 살펴보자.

어떠한가? N번째 발견마다 구간합이 0이 되는 구간이 N-1개씩 생겨난다.

그러면, 우리는 위의 누적합 아이디어를 이용해서 B의 원소들 중 누적합이 0이 되는 구간들을 찾아내면 되는 것이다. 이전 예시로 돌아가서, 누적합이 0이 되는 부분수열들을 찾아보자.
(다시 강조하지만 누적합이 0이 되는 구간들이란, 이전에 나왔던 누적합이 또 나오는 경우를 말한다.)

어떠한 누적합이 몇 번 나왔는지를 기록하기 위해서 collections 모듈의 defaultdict를 사용하자.
리스트 A를 순회하면서 누적합을 업데이트시켜주고, answer에 defaultdict에 구한 누적합에 대한 카운트를 더한 후, defaultdict에서 해당 누적합의 카운트를 1씩 증가시켜준다.

이때, defaultdict[0]은 다른 것과 다르게 먼저 1로 초기화해주는데, 그 이유는 아무것도 더하지 않았을 때에도 합이 0이기 때문이다.

코드

from collections import defaultdict

def solutions(A, S):
	B = [a - S for a in A]
    mem = defaultdict(int)
    mem[0] = 1
    
    answer, prefix_sum = 0, 0
    for b in B:
    	prefix_sum += b
        ans += mem[prefix_sum]
        mem[prefix_sum] += 1
    
    return answer