Prefix Sum

문제

---

1 부터 10 까지의 숫자들이 담겨져있는 배열 arr 이 있다고 가정한다.

여기서, 특정 인덱스 arr[a] 부터 arr[b] 까지의 누적합을 알아내고 싶다.

이 연산으로 딱 한번이 아닌, 여러번 연속으로 알아내고 싶다.

이럴 땐 어떻게 해야하는가?

for 문을 이용한 누적합

---

int main()
{
    int n; // [수의 개수] 1 ~ 10까지 담을테니, 10으로 입력한다.
    int m; // [출력 개수] 테스트케이스
    int a; // vec[a] 부터
    int b; // vec[b] 까지

    cin >> n >> m;

    vector<int> vec(n + 1); // 0 번자리를 제외하고, 인덱스 번호에 맞게 채우기위해 1부터 시작.

    for (int i = 1; i <= n; i++)
    {
        cin >> vec[i]; // 각 인덱스 번호에 맞게 1 ~ 10까지 담는다.
    }


    for (int i = 0; i < m; i++) // 테스트케이스
    {
        int sum = 0;

        cin >> a >> b;

        for (int j = a; j <= b; j++) // vec[a] 부터 vec[b] 까지 더하기
        {
            sum += vec[j];
        }
        cout << sum << '\n'; // 출력 후 다음 케이스 실행
    }
    return 0;
}

누적합 알고리즘을 모른다면 이렇게 푸는 것이 일반적이다.

첫번째 for 문으로 몇번을 시행할 건지, 두번째 for 문으로 입력한 구간 사이를 매번 더해줄 것이다.

이렇게 되면, 테스트케이스의 수, 혹은 인덱스의 범위가 늘어날 수록, 그리고 매번 순회하며 더해줄테니, 시간이 오래걸린다.

누적합 알고리즘

---

int main()
{
    int arr[11], psum[11];

    int n; // [수의 개수] 1 ~ 10까지 담을테니, 10으로 입력한다.
    int m; // [출력 개수] 테스트케이스
    int a; // arr[a] 부터
    int b; // arr[b] 까지

    cin >> n >> m;

    for (int i = 1; i <= n; i++) // make psum logic
    {
        cin >> arr[i];
        psum[i] = psum[i - 1] + arr[i];
    }

    for (int i = 0; i < m; i++) // query
    {
        cin >> a >> b;
        cout << psum[b] - psum[a - 1] << '\n';
    }

    return 0;
}

만약 arr 에 아래와 같이 숫자를 입력했다면, psum 배열엔 이렇게 담긴다.

arr = [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10]
psum[0] = 0 (초기값, 0번 인덱스는 사용하지 않음)
psum[1] = psum[0] + arr[1] = 0 + 1 = 1
psum[2] = psum[1] + arr[2] = 1 + 2 = 3
psum[3] = psum[2] + arr[3] = 3 + 3 = 6
psum[4] = psum[3] + arr[4] = 6 + 4 = 10
psum[5] = psum[4] + arr[5] = 10 + 5 = 15
psum[6] = psum[5] + arr[6] = 15 + 6 = 21
psum[7] = psum[6] + arr[7] = 21 + 7 = 28
psum[8] = psum[7] + arr[8] = 28 + 8 = 36
psum[9] = psum[8] + arr[9] = 36 + 9 = 45
psum[10] = psum[9] + arr[10] = 45 + 10 = 55

현재 psum 에는 이런식으로 값이 할당되어 있는데, 그림으로 표현하면 쉽게 이해할 수 있다.

이번엔 위 코드의에서 구간쿼리를 진행하는 부분이다.

    for (int i = 0; i < m; i++) // query
    {
        cin >> a >> b;
        cout << psum[b] - psum[a - 1] << '\n';
    }

만약 숫자 4 부터 7 까지의 누적합을 알고싶다고 하면, 아래와 같이 진행된다.

이렇게 하면, 각 쿼리에 대한 시간복잡도는 O(1) 상수시간이 된다.