구간 업데이트와 구간 합이 모두 O(logN)에 가능한 펜윅트리 (Fenwick Tree Lazy Propagation)

(2019.08.11 수정)

요즘 이 글의 검색량이 많아져서 조금 수정을 했다.

일반 펜윅트리(Fenwick Tree)에 관한 설명은 어디에나 많이 찾아볼 수 있다.

오늘은 구간 업데이트(Range Update)와 구간 합(Range Sum)이 O(lgN)에 가능한 특별한 Fenwick Tree에 대해서 포스팅해볼 것이다.

제목에 Fenwick Tree Lazy Propagation이라 작성했지만

실제로 Segment Tree Lazy Propagation처럼 update시에 점만 찍어뒀다가 query시에 한꺼번에 update하는 방식은 아니다.

다만 Segment Tree Lazy Propagation처럼 구간합과 구간쿼리가 모두 O(logN)에 가능하다는 뜻으로 썼다.

일반 펜윅트리: https://www.acmicpc.net/blog/view/21 (설명은 여기가 제일 깔끔하게 잘 되어있다.)

모든 알고리즘이나 자료구조가 그렇지만, 아는 것과 활용하는 것은 별개다.

특히나 이 펜윅트리는 사용하기에 따라서 정말 많은 문제들을 해결할 수 있다.

일반적으로 펜윅트리는 어떤 구간[1,N]에 대해서 구간 합을 구하고자 하는데

이 구간의 값이 계속해서 갱신될 때 구간의 합을 O(lgN)에 얻기 위해 쓴다.

다만, 해당 구간의 값이 갱신될 때에도 O(lgN)만큼 시간이 든다.

펜윅트리에 대한 설명은 위의 BOJ 블로그에도 써있으므로

코드만 간단하게 제시하겠다

#include <cstdio>
#include <cassert>
using namespace std;
constexpr int N=100000;
int fen[N+1];
void update(int at, int by) {
    assert(1<=at && at<=N);
    for(;at<=N;at+=at&-at) fen[at]+=by;
}
int query(int at) {
    if(at==0) return 0;
 
    assert(1<=at && at<=N);
    int ret=0;
    for(;at;at-=at&-at) ret+=fen[at];
    return ret;
}
int main() {
    int Q; scanf("%d", &Q);
    while(Q--) {
        int type; scanf("%d", &type);
        if(type==1) { //update
            int at,by;
            scanf("%d%d",&at,&by);
            update(at,by);
        } else if(type==2) { //query
            int L,R;
            scanf("%d%d",&L,&R);
            printf("%d\n", query(R)-query(L-1));
        } else exit(1); //unreachable code
    }
    return 0;
}
 

위의 코드에서 봐야 할 부분은 void update(int, int) 함수의 for문 부분과

int query(int) 함수의 for문 부분이다.

위의 코드에서 Q는 앞으로 들어올 입력 줄의 개수이고,

type이 1이면 at위치에 by만큼 update를 하고,

type이 2이면 [L,R] 구간의 합을 출력한다.

위에서도 언급했듯이 query와 update의 시간복잡도는 O(lgN)이다.

그럼 위의 코드를 어떻게 개량해서 구간 업데이트(range update)를 할 수 있을까?

사실 구간 합을 O(lgN)에 구할 필요가 없다면 구간 업데이트는 따로 고민할 필요가 없다.

펜윅트리 자체가 구간에 대해 update를 해주고 있기 때문에 이를 이용하면 된다.

펜윅트리에 update(5,1)을 한다는 것은 사실 [5,N]에 모두 1씩 더해주는 것과 같다.

그럼 [5,N]구간 안에 있는 at을 query를 통해 확인해본다면 at에 1만큼 더해져 있는것을 확인할 수 있다.

위에서 설명했듯이,

일반적인 펜윅트리는 구간합이나 구간업데이트 둘 중 하나를 O(lgN)에 가능하게 해준다.

구간 합을 O(lgN)에 구하기 위해서는 구간 업데이트를 O(lgN)에 하는 것은 불가능하고,

반대로 구간 업데이트를 O(lgN)에 한다면 구간 합을 O(lgN)에 구하는 것도 불가능하다.

그럼 둘다 O(lgN)에 하기위해선 어떻게 해야할까??

이 방법은 예전에 유성(Meteors: https://www.acmicpc.net/problem/8217) 문제를 풀면서 알게 되었는데,

사실 이 문제에서 구간 업데이트와 구간 쿼리(구간합)가 동시에 필요하진 않다.

다만, 그 당시 나는 잘 몰랐기 때문에

세그먼트 트리(Segment Tree) 말고도 구현이 간단한 Fenwick Tree로 구간 업데이트를 O(lgN)에 하는 방법을 찾다가

우연히 발견하게 되었다. (다시 말하지만 구간 업데이트만 O(lgN)에 하려면 일반적인 펜윅트리로 해결하면 된다.)

아무튼 이 방법은 Petr가 발견한 것 같다.

이 내용은 그의 블로그에서 배운것을 작성한 것이다.

구간 업데이트와 구간 합을 O(lgN)에 구하는 코드는 아래와 같다.

#include <cstdio>
#include <cassert>
using namespace std;
using ll = long long;
constexpr int N=100000;
ll tadd[N+1], tmul[N+1];
void inupdate(int at, ll mul, ll add) {
    assert(1<=at && at<=N);
    for(;at<=N;at+=at&-at) {
        tmul[at] += mul;
        tadd[at] += add;
    }
}
void range_update(int l, int r, ll by) {
    inupdate(l, by, -by*(l-1));
    inupdate(r+1, -by, by*r);
}
ll range_query(const int at) {
    if(at==0) return 0;
    assert(1<=at && at<=N);
 
    ll mul=0, add=0;
    for(int i=at;i;i-=i&-i) {
        mul += tmul[i];
        add += tadd[i];
    }
    return at*mul+add;
}
int main() {
    int Q; scanf("%d", &Q);
    while(Q--) {
        int type; scanf("%d", &type);
        if(type==1) { //range_update
            int L,R,by;
            scanf("%d%d%d",&L,&R,&by);
            range_update(L,R,by);
        } else if(type==2) { //range_query
            int L,R;
            scanf("%d%d",&L,&R);
            printf("%lld\n", range_query(R)-range_query(L-1));
        } else exit(1); //unreachable code
    }
    return 0;
}
 

코드에 대한 설명은 아래와 같다.

위의 그림과 같이 [l,r] 구간을 by만큼 증가시킨다고 가정하자.

그럼 폐구간[l,r] 범위에 있는 at에서 구간합 sum[1,at] 값이 얼마일까?

by*at - by*(l-1)이다.

이를 그림으로 표현하면 아래와 같다.

그럼 여기서 at 위치에서 바로 구할 수 있는 값은 무엇일까?

다시말해서 sum[1,at]을 구하려고 할 때 내가 at이라는 입력을 넣으면

at*by - (l-1)*by 값을 얻고 싶다.

어떻게 해야할까??

일단 at값은 우리가 입력으로 넣어줄 것이므로 pass하고 나면,

at에 곱해지는 by와 -(l-1)*by가 남는다.

그런데 [l,r]구간이 아닌 [l',r']구간에서 by'만큼 구간 업데이트가 일어난다면,

그리고 폐구간[l',r'] 범위에 at이 있다면,

sum[1,at] = at*(by+by') - (l-1)*by - (l'-1)*by'이 된다.

이를 다시 정리해보면,

sum[1,at] = at*(by+by') + [{-(l-1)*by} + {-(l'-1)*by'}]

그러므로 위에서 가정한 구간 업데이트가 일어날 때,

by, by'에 해당하는 부분은 tmul이라는 변수에 저장해두고,

-(l-1)*by, -(l'-1)*by'에 해당하는 부분은 tadd라는 변수에 저장해두면 된다.

그런데 [l,r]구간의 모든 tmul값을 update 해야 한다.

그런데 이 방법을 O(r-l+1)이 아닌 O(log(r-l+1))에 구현하는 방법을 우리는 이미 알고 있다.

이것이 바로 일반 Fenwick Tree이다.

그리고 위에서는 [l,r], [l',r'] 구간에 속하는 부분만 고려했는데,

(r,INF), (r',INF)에 속하는 부분에는 [l,r]과 [l',r']에 업데이트 했던 방식과는 반대로 update를 진행하면 된다.

이를 그림으로 나타내면 아래와 같다.

[2019.07.22 수정]

** 위 그림에서 [1,at']의 구간 합이 잘못됐었다. (지금은 올바르게 수정한 상태)

수정 전 그림 자료: at'*(by-by) + r*by

수정 후 그림 자료: at'*(by-by) + r*by - (l-1)*by

그러므로 폐구간[l,r]을 by만큼 구간 업데이트 할 땐,

tmul[l,N] += by

tadd[l,N] += -(l-1)*by

tmul[r+1,N] += -by

tadd[r+1,N] += r*by

위와 같이 업데이트를 하면 된다.

나중에 sum[l,r]을 가져다 쓸 때는

sum[1,r] = tmul[1,r]*r + add[1,r]

sum[1,l-1] = tmul[1,l-1]*(l-1) + add[1,l-1]

=> sum[1,r] = sum[1,r] - sum[1,l-1]

위와 같은 방식으로 가져다 쓰면 된다.

참고로

업데이트 하는거

at+=at&-at으로 하는데,

이거 at-=at&-at으로 해도 된다.

이렇게 하면 update를 [1,at]에다가 by만큼 더해졌음을 체크하는 꼴..

그럼 query날려서 확인할 때도 반대로 해주면 된다.

다만, 이렇게 query를 통해 구간합을 구하면 (그러니까 query를 at+=at&-at으로 하면)

구간[at,INF]에 얼마만큼 더해졌는지 확인하는 것이 된다.

지금 위에서 설명한 내용이 무슨말인지 안다면, 아래의 요약까지도 다 이해할 수 있을 것이다.

1. 구간합만 O(lgN)에 구하고 싶으면 update를 at+=at&-at으로 해주고,

2. 구간업데이트만 O(lgN)에 하고 싶으면 update를 at-=at&-at으로 해주고,

3. 구간합과 구간업데이트 모두 O(lgN)에 하고 싶으면 내가 바로 위에 포스팅한 것 참고하면 된다.

이 글은 바로 위의 3번 내용을 포스팅 하기 위해서 쓴 것이고, 해당 내용은 Petr 블로그에서 1년 전쯤에 보고 배운것이다.

https://petr-mitrichev.blogspot.com/2013/05/fenwick-tree-range-updates.html

@ 2019.08.11

관련된 논문도 찾았다.

https://arxiv.org/abs/1311.6093

해당 논문은 Petr가 제시한 Fenwick Tree보다 조금 더 일반적인 경우에 대해 말하고 있지만 Petr가 블로그에 반년정도 먼저 posting을 하긴 했다.

다차원 배열에서 Range Update와 Range Query가 모두 O(4^d log^d N)에 되는 Fenwick Tree에 대해서 말하고 있다.