누적 합 Prefix Sum

 어떤 배열 arr에 대해 prefix sum 배열을 만들어 배열의 index i부터 index j까지의 합을 O(1)으로 풀 수 있는 알고리즘. 배열의 크기가 $O(n^{k})$일 때 초기화는 $O(n^{k})$, 구간 합을 구할 때는 O(1)의 시간이 걸린다.

 

 특정 구간의 합을 구하는 것이기 때문에 해당 구간의 값이 변화하지 않는 것을 통해, 특정 구간 안에 원소가 있는지 없는지도 쉽게 계산할 수 있다.(2차원도 적용 가능)

range sum을 구할 때 index를 주의하도록 하자.

 

1. 1차원 Prefix Sum : 초기화 $O(n)$, 쿼리당 O(1)

 range sum을 구할 때, start가 0이면 -1로 out of bound가 되어버리므로 start가 0일 때만 예외 처리를 해 주자.

// return partial sum
// psum[i] : sum from index 0 to index i
vector<int> setPsum(vector<int>& arr){
	vector<int> psum(arr.size());
	psum[0] = arr[0];

	for(int i = 1; i<arr.size(); i++){
		psum[i] = psum[i-1] + arr[i];
	}
	
	return psum;
}

// return range sum from index 'start' to index 'end'
int getRangeSum(vector<int>& psum, int start, int end){
	if(start == 0) return psum[end];
	return psum[end] - psum[start - 1];
}

int main(void) {
	...
	vector<int> psum = setPsum(arr);

	int start, end;
	cin>>start>>end;
	cout<<getRangeSum(psum, start-1, end-1)<<'\n';
    ...

	return 0;
}

 

 

2. 2차원 Prefix Sum : 초기화 $O(n^{2})$, 쿼리당 O(1)

 위 그림에서 좌상단을 (0, 0), 가운데를 (r1, c1), 우하단을 (r2, c2)라고 하면, (r1, c1)부터 (r2, c2)까지의 합은 [(전체) - (파랑 + 노랑) - (초록 + 노랑) + 노랑] 이다. 이를 수식으로 표현하자면

 

(r1, c1)부터 (r2, c2)까지 합 = psum[r2][c2] - psum[r1-1][c2] - psum[r2][c1 - 1] + psum[r1 - 1][c1 - 1]

 

만약 r1이 0이면 r1이 들어가는 항은 0, c1이 0이면 c1이 들어가는 항은 0이 된다. 그러나 index 0에 0을 추가하는 간단한 트릭을 이용해 초기화에 대한 코드 작성을 조금 편리하게 할 수 있다.

 

 초기화도 같은 맥락으로 psum[r][c] = value[r][c] + psum[r-1][c] + psum[r][c-1] - psum[r-1][c-1]로 하면 된다.

// get partial sum from point (r1, c1) to point (r2, c2)
int getRangeSum(vector<vector<int>>& psum, int r1, int c1, int r2, int c2){
	int result = psum[r2][c2];
	if(r1 != 0) result -= psum[r1 - 1][c2];
	if(c1 != 0) result -= psum[r2][c1 - 1];
	if(r1 != 0 && c1 != 0) result += psum[r1 - 1][c1 - 1];
	return result;
}

int main(void) {
	...	
	int N, M; cin>>N>>M;
	vector<vector<int>> board(N, vector<int>(N)), psum(N+1, vector<int>(N+1, 0));

	// input
	for(int i = 0; i<N; i++){
		for(int j = 0; j<N; j++){
			cin>>board[i][j];
		}
	}
	
    // initialize prefix sum
    // psum[r][c] : sum from (0, 0) to (r, c)
	for(int r = 1; r<=N; r++){
		for(int c = 1; c<=N; c++){
			psum[r][c] = board[r-1][c-1] + psum[r-1][c] + psum[r][c-1] - psum[r-1][c-1];
		}
	}

	int x1, y1, x2, y2;
	while(M--){
		cin>>x1>>y1>>x2>>y2;
		cout<<getRangeSum(psum, x1, y1, x2, y2)<<'\n';
	}

	return 0;
}