2차원 세그먼트 트리의 새로운 구현

2차원 세그먼트 트리는 세그먼트 트리의 각 노드들을 세그먼트 트리들로 구현하여 2차원 쿼리를 처리할 수 있는 자료구조이다. 이 글에서는 머지 소트 트리의 아이디어를 활용한 2차원 세그먼트 트리의 새로운 구현 방법에 대해 설명하고자 한다. 이 글은 세그먼트 트리와 머지 소트 트리에 대한 기본 지식이 있다는 가정 하에 작성되었으므로 이에 대한 이해가 부족하다면 잘 정리된 관련 자료들이 많으니 공부하고 오는 것을 추천한다. 

 

 

https://www.acmicpc.net/problem/15977

내가 위 문제를 풀면서 이 아이디어를 떠올리게 되었으므로 위 문제를 토대로 과정을 설명한다. 이 문제는 최종적으로 열을 선택하고 정렬하여 각 행이 모두 증가 수열이 되는 최대 열 개수를 구하는 문제로 변환된다. 행 개수가 2개 또는 3개인데, 2개일 때는 정렬 후 LIS를 구하는 문제로 간단히 해결되므로 행 개수가 3개일 때를 살펴 보자. 우선 생각해 볼 수 있는 것은 첫 번째 행을 기준으로 정렬하는 것이다. 이제 2가지 원소에 대해 LIS를 구하는 문제가 되었다. 

 

나는 이 문제를 세그 DP의 관점에서 접근하였다. 즉, 현재 원소를 $(x, y)$라 하면, $a < x$이고 $b < y$인 모든 $(a, b)$에 대해 $DP[(x,  y)] = max(DP[(a, b)]) + 1$이 될 것이다. 이런 구간에 대한 max값과 단일 원소에 대한 update는 세그먼트 트리로 해결할 수 있다. 

 

우선, x에 대한 대소 비교는 x를 좌표압축한 후 index로 사용하는 것으로 해결된다. 세그먼트 트리를 이용해 단일 원소에 대한 LIS를 푸는 방법과 같은 아이디어를 차용한 것이다. 이제 y를 처리하는 것이 문제인데... 여기서 생각한 것이 머지 소트 트리였다. x의 범위에 속하는 모든 y들을 저장할 수 있고 y들이 정렬되어 있어 x의 구간이 정해져 있을 때, y의 범위를 이분탐색을 이용해  결정하고 인덱스로 사용하는 것이 가능해진다.

 

y를 인덱스로 사용하는 것이 가능해진다? 여기서 내가 떠올린 방법이 세그먼트 트리를 만드는 것이다. 정확히는 머지 소트 트리의 깊이 별로 세그먼트 트리를 만드는 것이다. 머지 소트 트리에서 깊이 별로 y가 어떻게 정렬되어 있는지 생각해 보자. x의 구간들이 균등하게 분리되어 있고, 각 구간 별로 y가 정렬되어 있다. 머지 소트 트리에서 어떤 노드에 들어온 순간, x의 범위는 더 이상 고려할 필요가 없기 때문에 각 구간 별로 y가 커지는 순으로 index를 배정하면 y에 대해서도 해결하는 것이 가능하다. 다양한 구현 방법이 있겠지만, 나는 깊이 별로 세그먼트 트리를 만들고, 머지 소트 트리는 세그먼트 트리와 범위가 같으므로 배열을 log개 만들고 머지 소트 트리와 같은 방식으로 정렬하는 것으로 구현하였다. 이는 나중에 다시 설명하겠다. 결과적으로 log개의 세그먼트 트리를 만들게 되고 공간복잡도는 $O(NlogN)$, 시간복잡도는 $O(N(logN)^2)$이 되어 문제를 해결할 수 있다.

 

예시를 보면 이해가 쉬울 것이다.

위와 같은 배열이 있다고 하자. x는 인덱스로 사용될 것이기 때문에 인덱스를 x로, 값을 y로 생각하면 이해하기 편리하다. 세그먼트 트리를 2^N 크기로 구현할 것이기 때문에 편의를 위해 길이를 2의 제곱수인 8로 설정하였다.

각각의 박스를 하나의 노드로 생각하면, 머지 소트 트리는 위와 같이 시각화할 수 있다. 위 그림에서 아까 말한 방법으로 y에 인덱스를 배정하면 다음과 같다. 

이제 깊이 별로 세그먼트 트리를 만든다는 것이 어떤 의미인지 잘 전달되었을 것이다. 각각의 초록색 박스를 노드로 생각하면, 노드 안에서 y는 정렬되어 있으므로 y에 대해 이분탐색하면 해당하는 y의 인덱스를 구할 수 있다. 따라서 깊이 별 세그먼트 트리 내에서 인덱스 충돌 없이 범위에 대한 쿼리를 수행할 수 있다.

 

마지막으로, 구현 방법을 설명한다.

vector<tll> a(n);
for (ll j = 0; j < m; j++) {
    for (ll i = 0; i < n; i++) {
        cin >> a[i][j];
    }
}

sort(a.begin(), a.end());
vector<ll> v(n);
for (ll i = 0; i < n; i++) {
    v[i] = a[i][1];
}

sort(v.begin(), v.end());
v.erase(unique(v.begin(), v.end()), v.end());

 

여기까지는 기본적이다. 첫 번째 원소에 대해 정렬하여 두 번째, 세 번째 원소에 대한 LIS로 문제를 변환한고, x를 좌표 압축한다.

 

auto index = [&](ll val) {
    return lower_bound(v.begin(), v.end(), val) - v.begin();
};

vector<pll> b(n);
for (ll i = 0; i < n; i++) {
    a[i][1] = index(a[i][1]);
    b[i] = {a[i][1], a[i][2]};
}
sort(b.begin(), b.end());
for (ll d = 0; d < 19; d++) {
    ll t = 1 << (18 - d);
    for (ll i = 0; i < n; i++) {
        sv[d].emplace_back(b[i].second);
    }
    for (ll i = 0; i < n; i += t) {
        sort(sv[d].begin() + i, sv[d].begin() + min(n, i + t));
    }
}

 

이 부분이 내 아이디어의 핵심이다. x의 범위를 유지하기 위해 x, y를 배열에 pair로 저장한 후, 정렬한다. N이 최대 20,000이므로 머지 소트 트리의 최대 깊이는 18이 된다. 각 깊이 별로 배열을 만들고, 하나의 노드가 포함하는 구간의 길이에 해당하는 범위를 각각 정렬하면 머지 소트 트리 없이 머지 소트 트리와 같은 기능을 한다. 사실 머지 소트 트리로 구현해도 되지만, 구현의 편의성과 시간 단축을 위해 위와 같이 구현하였다.

 

struct max_seg {
    array<ll, MAX << 1> tree{};

    void update(ll idx, ll val) {
        tree[idx |= MAX] = val;
        while (idx >>= 1) {
            tree[idx] = max(tree[idx << 1], tree[idx << 1 | 1]);
        }
    }

    ll query(ll s, ll e) {
        ll ret = 0;
        for (s |= MAX, e |= MAX; s < e; s >>= 1, e >>= 1) {
            if (s & 1) ret = max(ret, tree[s++]);
            if (e & 1) ret = max(ret, tree[--e]);
        }
        return ret;
    }
} msg[19];

 

세그먼트 트리는 여러 개 깊이 별로 만들어야 하기 때문에 구조체로 구현하였다. query 함수에서 s는 닫힌 구간, e는 열린 구간으로 구현하였으니 유의하자.

 

void update(ll x, ll y, ll v, ll l = 0, ll r = MAX, ll cur = 1) {
    ll d = 31 - __builtin_clz(cur);
    ll it = lower_bound(sv[d].begin() + l, sv[d].begin() + min<ll>(r, sv[d].size()), y) - sv[d].begin();
    msg[d].update(it, v);
    if (r - l == 1) return;
    ll m = (l + r) >> 1;
    if (x < m) {
        update(x, y, v, l, m, cur << 1);
    } else {
        update(x, y, v, m, r, cur << 1 | 1);
    }
}

ll query(ll s, ll e, ll y, ll l = 0, ll r = MAX, ll cur = 1) {
    if (r <= s or e <= l) {
        return 0;
    } else if (s <= l and r <= e) {
        ll d = 31 - __builtin_clz(cur);
        ll it = lower_bound(sv[d].begin() + l, sv[d].begin() + r, y) - sv[d].begin();
        return msg[d].query(l, it);
    } else {
        ll m = (l + r) >> 1;
        ll r1 = query(s, e, y, l, m, cur << 1);
        ll r2 = query(s, e, y, m, r, cur << 1 | 1);
        return max(r1, r2);
    }
}

 

이제 마지막 부분이다. 깊이를 비트 연산을 이용해 구한 후, 아까 구현한 머지 소트 트리와 같이 정렬된 배열들에서 현재 깊이에 해당하는 배열을 본다. 헷갈리지 않도록 첨언하면 query 함수의 update 함수에서 s, e, l, r은 모두 x를 좌표압축한 인덱스이다. 현재 노드가 해당하는 구간이 우리가 보고 싶은 구간 내에 완벽히 들어 맞는다면 x의 범위는 더 이상 고려할 필요가 없다. 아까 구해 둔 현재 구간 내의 y좌표에서 이분탐색을 하고, 현재 깊이의 세그먼트 트리에 쿼리를 날리면 구현이 끝난다. 

 

vector<ll> dp(n);
for (ll i = 0; i < n; i++) {
    dp[i] = query(0, a[i][1], a[i][2]) + 1;
    update(a[i][1], a[i][2], dp[i]);
}
cout << *max_element(dp.begin(), dp.end());

 

지금까지 구현한 자료구조를 이용해 dp를 수행하면 정답을 구할 수 있다.

 

 

사실 이게 2차원 세그먼트 트리가 맞는지도 잘 모르겠다. 보통은 각 세그먼트 트리의 노드 내에 또 하나의 세그먼트 트리를 만드는 것으로 구현하는 듯 하다. 다른 분들의 구현도 참고해 보았지만 나는 내 방식의 구현이 가장 간단하다고 느꼈다. 세그먼트 트리를 log개만 만들고 관리하면 되기 때문에 세그먼트 트리를 노드의 개수만큼 만드는 것보다는 구현이 간편할 수밖에 없다. 그나마 머지 소트 트리의 아이디어를 활용하는 부분이 좀 난해하긴 한데 이런 방식으로 구현한 코드를 지금까지 보지 못했기 때문에 나만의 구현이 되지 않을까 싶다. 내가 2차원 세그먼트 트리를 알았더라면 이 고생을 안해도 되었겠지만, 모르던 자료구조를 자력으로 생각해서 풀었다는 데에 의의를 두려고 한다.