2D Segment Tree

2D Segment Tree는 y축에 Segment Tree 한개, x축에 Segment Tree 한개를 이용하여 이차원 좌표상의 한 점 업데이트 및 구간 연산 쿼리를 빠르게 처리할 수 있는 자료구조이다.

 

y축 Segment Tree는 각 노드마다 x축 Segment Tree 를 하나씩 갖고 있으며, 정확히는 구간 [yl, yr]을 담당하고 있을 때, 각 노드가 [yl, yr][x, x] 를 노드로 갖는 x축 Segment Tree 를 의미한다.

 

그냥 구현하면 공간 복잡도가 $O(N^2)$ 이니, Dynamic Segment Tree로 X축, Y축을 모두 구현할 필요가 있다. Dynamic Segment Tree를 사용하다 보니, Pointer 을 사용한 구현과 벡터 기반의 구현 두가지가 모두 존재한다.

 

1. Pointer 을 사용한 구현

#include <bits/stdc++.h>
#include "game.h"
using namespace std;

typedef long long ll;
typedef pair<int, int> pii;
typedef pair<ll, ll> pll;

int r, c;

struct Node1
{
    Node1() : val(0), pos(-1), lc(NULL), rc(NULL) {}
    ll val, pos; Node1 *lc, *rc;
};

struct Node2
{
    Node2() : val(NULL), lc(NULL), rc(NULL) {}
    Node1 *val; Node2 *lc, *rc;
};

Node2 *root;

void init(int R, int C)
{
    r=R; c=C;
    root=new Node2();
}

void update1(Node1* node, int tl, int tr, int x, ll val)
{
    if(tl==tr) { node->val=val; return; }
    int mid=tl+tr>>1;

    if(node->pos!=-1)
    {
        if(node->pos<=mid)
        {
            node->lc=new Node1();
            node->lc->pos=node->pos;
            node->lc->val=node->val;
        }
        else
        {
            node->rc=new Node1();
            node->rc->val=node->val;
            node->rc->pos=node->pos;
        }
        node->pos=-1;
    }

    if(x<=mid)
    {
        if(node->lc==NULL)
        {
            node->lc=new Node1();
            node->lc->pos=x;
            node->lc->val=val;
        }
        else update1(node->lc, tl, mid, x, val);
    }
    else
    {
        if(node->rc==NULL)
        {
            node->rc=new Node1();
            node->rc->pos=x;
            node->rc->val=val;
        }
        else update1(node->rc, mid+1, tr, x, val);
    }
    node->val=0;
    if(node->lc!=NULL) node->val=__gcd(node->val, node->lc->val);
    if(node->rc!=NULL) node->val=__gcd(node->val, node->rc->val);
}

ll query1(Node1* node, int tl, int tr, int lx, int rx)
{
    if(tr<lx || rx<tl) return 0;
    if(lx<=tl && tr<=rx) return node->val;
    if(node->pos!=-1)
    {
        if(lx<=node->pos && node->pos<=rx) return node->val;
        else return 0;
    }
    int mid=tl+tr>>1;
    ll ret=0;
    if(node->lc!=NULL) ret=__gcd(ret, query1(node->lc, tl, mid, lx, rx));
    if(node->rc!=NULL) ret=__gcd(ret, query1(node->rc, mid+1, tr, lx, rx));
    return ret;
}

void update2(Node2* node, int tl, int tr, int y, int x, ll val)
{
    if(node->val==NULL) node->val=new Node1();
    if(tl==tr) { update1(node->val, 1, c, x, val); return; }
    int mid=tl+tr>>1;
    if(y<=mid)
    {
        if(node->lc==NULL) node->lc=new Node2();
        update2(node->lc, tl, mid, y, x, val);
    }
    else
    {
        if(node->rc==NULL) node->rc=new Node2();
        update2(node->rc, mid+1, tr, y, x, val);
    }
    ll t=0;
    if(node->lc!=NULL) t=__gcd(t, query1(node->lc->val, 1, c, x, x));
    if(node->rc!=NULL) t=__gcd(t, query1(node->rc->val, 1, c, x, x));
    update1(node->val, 1, c, x, t);
}

ll query2(Node2 *node, int tl, int tr, int ly, int ry, int lx, int rx)
{
    if(tr<ly || ry<tl) return 0;
    if(ly<=tl && tr<=ry)
    {
        if(node->val!=NULL) return query1(node->val, 1, c, lx, rx);
        else return 0;
    }
    int mid=tl+tr>>1;
    ll ret=0;
    if(node->lc!=NULL) ret=__gcd(ret, query2(node->lc, tl, mid, ly, ry, lx, rx));
    if(node->rc!=NULL) ret=__gcd(ret, query2(node->rc, mid+1, tr, ly, ry, lx, rx));
    return ret;
}

void update(int P, int Q, ll K)
{
    P++; Q++;
    update2(root, 1, r, P, Q, K);
}

ll calculate(int P, int Q, int U, int V)
{
    P++; Q++; U++; V++;
    return query2(root, 1, r, P, U, Q, V);
}

 

2. 벡터 기반의 구현

#include "game.h"
#include <bits/stdc++.h>
using namespace std;

typedef long long ll;
typedef pair<int, int> pii;
typedef pair<ll, ll> pll;

int R, C;

struct Node1
{
    ll val; int lc, rc, pos;
    Node1() : val(0), lc(-1), rc(-1), pos(-1) {}
};
vector<Node1> nodes1;
int newNode1()
{
    nodes1.push_back(Node1());
    return (int)nodes1.size()-1;
}

struct Node2
{
    int val; int lc, rc;
    Node2() : val(-1), lc(-1), rc(-1) {}
};
vector<Node2> nodes2;
int newNode2()
{
    nodes2.push_back(Node2());
    return (int)nodes2.size()-1;
}

void update1(int node, int tl, int tr, int x, ll val)
{
    //printf("%d %d %d %d %lld\n", node, tl, tr, x, val);
    if(tl==tr)
    {
        nodes1[node].val=val;
        return;
    }

    int mid=tl+tr>>1;
    if(nodes1[node].pos!=-1)
    {
        if(nodes1[node].pos<=mid)
        {
            int t=newNode1(); nodes1[node].lc=t;
            nodes1[nodes1[node].lc].pos=nodes1[node].pos;
            nodes1[nodes1[node].lc].val=nodes1[node].val;
        }
        else
        {
            int t=newNode1(); nodes1[node].rc=t;
            nodes1[nodes1[node].rc].pos=nodes1[node].pos;
            nodes1[nodes1[node].rc].val=nodes1[node].val;
        }
        nodes1[node].pos=-1;
    }

    if(x<=mid)
    {
        if(nodes1[node].lc==-1)
        {
            int t=newNode1(); nodes1[node].lc=t;
            nodes1[nodes1[node].lc].pos=x;
            nodes1[nodes1[node].lc].val=val;
        }
        else update1(nodes1[node].lc, tl, mid, x, val);
    }
    else
    {
        if(nodes1[node].rc==-1)
        {
            int t=newNode1(); nodes1[node].rc=t;
            nodes1[nodes1[node].rc].pos=x;
            nodes1[nodes1[node].rc].val=val;
        }
        else update1(nodes1[node].rc, mid+1, tr, x, val);
    }
    ll t=0;
    if(nodes1[node].lc!=-1) t=__gcd(t, nodes1[nodes1[node].lc].val);
    if(nodes1[node].rc!=-1) t=__gcd(t, nodes1[nodes1[node].rc].val);
    nodes1[node].val=t;
}

ll query1(int node, int tl, int tr, int xl, int xr)
{
    //printf("%d %d %d %d %d\n", node, tl, tr, xl, xr);
    if(xr<tl || tr<xl) return 0;
    if(xl<=tl && tr<=xr) return nodes1[node].val;
    if(nodes1[node].pos!=-1)
    {
        if(xl<=nodes1[node].pos && nodes1[node].pos<=xr) return nodes1[node].val;
        else return 0;
    }
    int mid=tl+tr>>1;
    ll ret=0;
    if(nodes1[node].lc!=-1) ret=__gcd(ret, query1(nodes1[node].lc, tl, mid, xl, xr));
    if(nodes1[node].rc!=-1) ret=__gcd(ret, query1(nodes1[node].rc, mid+1, tr, xl, xr));
    return ret;
}

void update2(int node, int tl, int tr, int y, int x, ll val)
{
    //printf("%d %d %d %d %d %lld\n", node, tl, tr, y, x, val);
    if(nodes2[node].val==-1) { int t=newNode1(); nodes2[node].val=t; }
    if(tl==tr)
    {
        update1(nodes2[node].val, 1, C, x, val);
        return;
    }
    int mid=tl+tr>>1;
    if(y<=mid)
    {
        if(nodes2[node].lc==-1) { int t=newNode2(); nodes2[node].lc=t; }
        update2(nodes2[node].lc, tl, mid, y, x, val);
    }
    else
    {
        if(nodes2[node].rc==-1) { int t=newNode2(); nodes2[node].rc=t; }
        update2(nodes2[node].rc, mid+1, tr, y, x, val);
    }
    ll t=0;
    if(nodes2[node].lc!=-1) t=__gcd(t, query1(nodes2[nodes2[node].lc].val, 1, C, x, x));
    if(nodes2[node].rc!=-1) t=__gcd(t, query1(nodes2[nodes2[node].rc].val, 1, C, x, x));
    update1(nodes2[node].val, 1, C, x, t);
}

ll query2(int node, int tl, int tr, int yl, int yr, int xl, int xr)
{
    //printf("%d %d %d %d %d %d %d\n", node, tl, tr, yl, yr, xl, xr);
    if(yr<tl || tr<yl) return 0;
    if(yl<=tl && tr<=yr)
    {
        if(nodes2[node].val!=-1) return query1(nodes2[node].val, 1, C, xl, xr);
        else return 0;
    }
    int mid=tl+tr>>1;
    ll ret=0;
    if(nodes2[node].lc!=-1) ret=__gcd(ret, query2(nodes2[node].lc, tl, mid, yl, yr, xl, xr));
    if(nodes2[node].rc!=-1) ret=__gcd(ret, query2(nodes2[node].rc, mid+1, tr, yl, yr, xl, xr));
    return ret;
}

int root;

void init(int _R, int _C)
{
    int i, j;
    R=_R; C=_C;
    root=newNode2();
}

void update(int P, int Q, ll K)
{
    int i, j;
    P++; Q++;
    update2(root, 1, R, P, Q, K);
}

ll calculate(int P, int Q, int U, int V)
{
    int i, j;
    P++; Q++; U++; V++;
    return query2(root, 1, R, P, U, Q, V);
}

 

위 코드는 IOI 2013 Game 의 풀이이다.

 

2D Segment Tree는 공간, 시간 복잡도 면에서 매우 무거운 자료구조이기 때문에, 몇몇 최적화를 해줄 수 있다.

그중 하나는 "경로 압축"으로, 다이나믹 세그트리에서도 사용될 수 있다.

 

[l, r]구간에 pos가 새로 들어와, 어떤 새로운 노드를 추가했다면, 이 노드를 원래대로라면 리프까지 쭉 늘여뜨리게 된다. 하지만 이 긴 경로들은 모두 필요 없고, lazy propagation 과 비슷한 느낌으로 위치를 표시해 놓고 바로 위에 달아 놓고, 필요할 때마다 내려주면 된다.

 

시간 복잡도 : $O(Qlog^2{N})$

공간 복잡도 : $O(Qlog^2{N})$