프로그래밍 입문 수업 - 구간합
저번에 특정 구간의 소수를 구하는 것을 설명드렸는데, 이번엔 어떤 배열에서 특정 구간의 합을 구하는 문제입니다.
이 알고리즘은 저번보다 훨씬 간단하니 금방 익힐 수 있고, 최근 기업 코딩 테스트에서 소수 구하는 유형보다 훨씬 많이 나왔어요! 👍
그럼 살펴봅시다!
아래와 같은 배열이 주어지고
| 3 | 2 | 1 | 4 | 5 | 7 | 2 | 8 | 0 | 10 |
이 배열의 특정 인덱스 i부터 j까지 (단, i<=j)의 합을 묻는 쿼리가 주어질 때 구간합을 묻는 문제라고 합니다.
이 문제는 간단히 (j-i+1)번 더하기 연산을 통해 계산해줄 수 있습니다!
예시 코드
#include <stdio.h>
int main()
{
int arr[10] = { 3, 2, 1, 4, 5, 7, 2, 8, 0, 10 };
int Q; // 쿼리의 개수
scanf("%d", &Q);
for (int i = 0; i < Q; i++)
{
int s, e; //start, end
scanf("%d %d", &s, &e);
int sum = 0;
for(int i=start; i<=end; i++){
sum += arr[i];
}
printf("%d\n", sum);
}
return 0;
}
하지만, 만약에 배열의 원소가 100만개이고 쿼리가 100만 개 정도 주어지면 무조건 문제가 발생하는데,
위 코드는 쿼리 한 개 당 최대 100만 번 연산을 하게되기 때문이죠 (Tip. 보통 1억 번 연산을 하면 1초 정도 소요돼서 시간 초과 가능성이 있다.)
따라서 우리는 구간합을 도와줄 수 있는 또 다른 배열을 하나 만들어 줄 건데, 이름은 psum이라고 해봅시다.
이 배열의 정의는 다음과 같다.
psum[i] = 0~i번째 인덱스까지의 원소 합
참고로 psum은 prefix sum의 줄임말입니다!
원래 배열과 비교해보면 다음과 같은 차이가 있습니다.
원래 배열(arr)
| 3 | 2 | 1 | 4 | 5 | 7 | 2 | 8 | 0 | 10 |
psum 배열
| 3 | 5 | 6 | 10 | 15 | 22 | 24 | 32 | 32 | 42 |
이제 이 배열의 활용법을 알아볼 건데, 예를 들어 3~5 구간의 합을 알기 위해선 두 가지 정보가 필요합니다.
원래 배열
| 3 | 2 | 1 | 4 | 5 | 7 | 2 | 8 | 0 | 10 |
| 3 | 2 | 1 | 4 | 5 | 7 | 2 | 8 | 0 | 10 |
눈치 빠르신 분들은 이미 아셨겠지만, 그대로 파란 부분 - 빨간 부분을 해주면 3~5 구간의 합이 나옵니다.
그런데, 이 부분들은 위에서 미리 구해줬습니다!
- 파란 부분(0~5) = psum[5]
- 빨간 부분(0~2) = psum[2]
따라서 psum[5] - psum[2] 만 해주면 바로 구간합을 구할 수 있는 것이죠!
별거 아닌 것처럼 보이지만 여러 번 하면 시간 차이는 엄청납니다 🤦♂️
요약하자면, start ~ end의 구간합은 (0~end)의 합 - (0~start-1)의 합으로 구할 수 있다는 점이 이 알고리즘의 핵심입니다!
만약 배열이 업데이트되는 조건이 추가되면, 세그먼트 트리라는 알고리즘이 필요한데 정신 건강에 해로우니 조심하세요!
(참고. 세그먼트 트리 : https://m.blog.naver.com/kks227/220791986409)
예시 코드
#include <stdio.h>
int main()
{
int arr[10] = { 3, 2, 1, 4, 5, 7, 2, 8, 0, 10 };
int psum[10] = {0,};
//누적합 구하는 방법
for(int i=0; i<10; i++){
if(i == 0) psum[i] = arr[i];
else psum[i] = psum[i-1] + arr[i];
}
int Q; // 쿼리의 개수
scanf("%d", &Q);
for (int i = 0; i < Q; i++)
{
int s, e; //start, end
scanf("%d %d", &s, &e);
printf("%d\n", psum[e] - psum[s-1]);
}
return 0;
}