Segment Tree Beats 15분만에 배우고 사용하기

사전 지식

Lazy Segment Tree의 구현

 

Atcoder Library의 Lazy Segment Tree 설명: https://atcoder.github.io/ac-library/production/document_en/lazysegtree.html

레이지 세그먼트 트리는 타입 $S$와 $S \rightarrow S$로 가는 함수들의 부분집합 $F$가 존재해서, 각 원소가 $S$ 타입인 배열에서 다음 일을 할 수 있습니다.

• 배열의 연속된 구간과 $f \in F$가 주어지면, 연속된 구간의 각 원소를 $s_i$에서 $f(s_i)$로 바꿉니다.
• 배열의 연속된 구간$[s_l, \cdots, s_r]$이 주어지면, $s_l \circ \cdots \circ s_r$ 을 구합니다.

여기서 연산 $\circ$와 함수의 집합 $F$는 다음 성질을 만족해야 합니다.

• 결합법칙: $a, b, c \in S$에 대해 $a \circ (b \circ c) = (a \circ b) \circ c$
• $F$가 결합에 대해 닫혀있음: $f, g \in F$ 이면 $f \circ g \in F$
• $F$가 연산을 보존함: $f \in F, a, b \in S$에 대해, $f(x) \circ f(y) = f (x \circ y)$

$S$를 문제에서 준 그대로 사용하면 해결하지 못할 수도 있을 수 있습니다. 이런 경우에는 $S$에 추가적인 정보를 같이 저장하면 됩니다. 예를 들어 연속된 구간에 특정 값을 더하는 연산을 해야하고 연속된 구간의 전체 합을 구해야 하는 경우에는 $S$가 합과 크기를 저장하고 있으면, 다음과 같이 정의하면 됩니다.

• $S$는 구간의 합과 구간에 있는 원소의 개수를 들고 있습니다.
• 배열의 각 원소 $v$를 $(v, 1)$로 표현합니다.
  • 원소 하나의 합은 $v$, 원소 하나의 개수는 $1$입니다.
• 두 구간의 연결은 $(sum_1, size_1) \circ (sum_2, size_2) = (sum_1+sum_2, size_1 + size_2)$로 표현합니다.
• 배열의 각 원소에 $a$를 더하는 연산은 $f_a((sum, size)) = (sum + a*size, size)$로 표현합니다.
  • 각 원소에 $a$를 더하면, 총합은 배열의 길이에 $a$를 곱한것 만큼 늘어납니다.

Segment Tree Beats

세그먼트 트리를 쓰다 보면 다음 상황에 맞닥뜨립니다.

• $F$를 적용하기 위해서, $S$에 너무 많은 정보를 저장해야 하는 경우가 있습니다.

이를 해결하기 위해서, 우리는 $S$에 있는 정보를 압축해서 저장합니다. 이 경우, $f(s)$가 유일하게 결정되지 않을 수도 있습니다. 이런 경우에 (자식 노드에 저장된 $s = s_l \circ s_r$에 대해) $f(s_l \circ s_r) = f(s_l) \circ f(s_r)$을 이용해서, 값을 다시 계산해줍니다.

 

다음과 같은 예를 들어봅시다: $f_a(x) = \min(a, x)$이고, $x \circ y = x+y$입니다. 이를 처리하기 위해서 $S$에 추가적인 정보를 저장하기 위해서는, 모든 배열을 다 저장해야 합니다. 하지만 $S$에 있는 정보를 압축해서 최댓값 $mx$, 최댓값의 개수 $cnt$, 두 번째 최댓값 $mx2$, 전체의 합 $sum$을 저장한다고 해 봅시다.

 

$f_a((mx, cnt, mx2, sum))$은 다음과 같이 작동합니다.

• $a \ge mx$이면, $f_a$를 적용하지 않아도 값이 같습니다. 즉, $f_a((mx, cnt, mx2, sum)) = (mx, cnt, mx2, sum)$입니다.
• $mx2 < a < mx$이면, $cnt$개의 $mx$원소만 $a$로 값이 바뀌고, 합은 $cnt \cdot (mx-a)$만큼 증가합니다. 즉, $f_a((mx, cnt, mx2, sum)) = (a, cnt, mx2, sum + cnt \cdot (mx-a))$입니다.
• $mx2 \le a$이면 $S$에 현재 있는 정보만 가지고 값을 계산할 수 없습니다. 자식 노드의 정보를 활용합니다.

앞의 두 경우는 일반 세그먼트 트리와 시간복잡도가 똑같습니다. 하지만 세 번째 경우가 일어나면 추가적으로 연산이 필요합니다. 쿼리를 처리하는 전체에서 사용하는 시간은 (세그먼트 트리에 드는 기본 시간) + (세 번째 경우가 일어나는 횟수)가 됩니다.

세 번째 경우가 자주 발생하지 않는다는 사실은, 세 번째 경우가 발생하면 노드들에 담긴 서로 다른 수의 개수가 적어도 $1$ 감소한다는 사실을 관찰을 할 수 있습니다. 최초에 노드에 담긴 서로 다른 수의 개수는 $O(n \log n)$이고 한 쿼리는 노드에 저장된 서로 다른 수의 개수를 최대 $O(\log n)$개 증가시키기 때문에, 세 번째 경우는 최대 $O((n+q) \log n)$번 일어나게 됩니다.

문제에 따라 적당히 $S$를 정해서 계산이 실패하는 횟수를 줄이면, 문제를 해결할 수 있습니다.

Atcoder Library

이는 Atcoder Library를 조금 바꿔서 해결할 수 있습니다.

우리는 $f(x)$의 계산이 실패하는 경우를 위해서, 기존의 mapping함수를 try_mapping함수로 바꿀 것 입니다. 이 함수는 bool try_mapping(F& f, S s, S& r);과 같은 함수 형태를 가졌으며, 계산이 성공하면 $f(s)$를 $r$에 저장하고 true를 반환합니다. 그렇지 않은 경우 false를 반환합니다.

 

그 후, 위에 설명된 점을 구현하기 위해 lazysegtree.hpp를 복사해서 segbeats.hpp를 만든 이후에, 다음 diff에 해당하는 변화를 적용합니다.

1,2c1,2
< #ifndef ATCODER_LAZYSEGTREE_HPP
< #define ATCODER_LAZYSEGTREE_HPP 1
---
> #ifndef ATCODER_SEGBEATS_HPP
> #define ATCODER_SEGBEATS_HPP 1
17c17
<           S (*mapping)(F, S),
---
>           bool (*try_mapping)(F&, S, S&),
79c79
<         d[p] = mapping(f, d[p]);
---
>         assert(try_mapping(f, d[p], d[p]));
177,178c177,182
<         d[k] = mapping(f, d[k]);
<         if (k < size) lz[k] = composition(f, lz[k]);
---
>         bool res = try_mapping(f, d[k], d[k]);
>         if (k < size)
>         {
>             lz[k] = composition(f, lz[k]);
>             if(!res) push(k), update(k);
>         } else assert(res);
189c193
< #endif  // ATCODER_LAZYSEGTREE_HPP
---
> #endif  // ATCODER_SEGBEATS_HPP

17번째 줄은 함수 형태를 바꿉니다. 79번째 줄은 원소 하나에 대한 apply를 적용할 때 항상 성공해야 한다는 의미입니다. 177번째 줄은 계산이 실패했을 경우 아래쪽 노드에 대한 정보를 가져와서 갱신해야 한다는 의미입니다.

 

다음은 위에서 설명한 쿼리를 해결하는 수열과 쿼리 26 문제를 위 방법으로 해결한 코드입니다.

#include <bits/stdc++.h>
#include <atcoder/segbeats>
using namespace std;

struct S {
    int mx, mxcnt, mx2; long long sum;
    S() : mx(0), mxcnt(0), mx2(0), sum(0) {}
    S(int v) : mx(v), mxcnt(1), mx2(0), sum(v) {}
};
S op(S a, S b) {
    if (a.mx < b.mx) swap(a, b);
    if (a.mx == b.mx) a.mxcnt += b.mxcnt, a.mx2 = max(a.mx2, b.mx2);
    else a.mx2 = max(b.mx, a.mx2);
    a.sum += b.sum;
    return a;
}
S e() { return S(); }
using F = int;
bool try_mapping(F& f, S s, S &r) {
    if (s.mx <= f) { r = s; return true; }
    else if (s.mx2 < f && f < s.mx) {
        s.sum -= (long long)(s.mx - f) * s.mxcnt;
        s.mx = min(s.mx, f);
        r = s; return true;
    } else return false;
}
F composition(F f, F g) { return min(f, g); }
F id() { return INT_MAX; }

using Seg = atcoder::segbeats<S, op, e, F, try_mapping, composition, id>;

int main() {
    ios::sync_with_stdio(false); cin.tie(0);
    int N; cin >> N;
    vector<S> V;
    for (int i = 0; i < N; i++) {
        int t; cin >> t;
        V.emplace_back(t);
    }
    Seg seg(V);

    int M; cin >> M;
    for (int i = 0; i < M; i++) {
        int t, l, r; cin >> t >> l >> r;
        if (t == 1) {
            int x; cin >> x;
            seg.apply(l - 1, r, x);
        } else {
            auto R = seg.prod(l - 1, r);
            cout << (t == 2 ? (long long)R.mx : R.sum) << '\n';
        }
    }
}

Why don't we go deeper...

수열과 쿼리 28을 풀어봅시다. 연속된 구간에 특정한 수를 더해야 하고, $f(x) = \lfloor \sqrt x \rfloor$를 적용해야 하며, 구간의 합을 구해야 합니다. 여기서의 $F$는 생각보다 잘 정의되지 않는다는것을 알 수 있습니다. 주어진 $F(x)$를 합성하면, 우리는 기존에 세그먼트 트리에서 맞딱뜨린 $S$의 경우처럼 $F$에 대한 정보를 너무 많이 저장해야합니다.

 

위에서 사용한 방법은 흔히 사용되는 Segment Tree Beats를 완전히 커버하지 못합니다. 이 세그먼트 트리는 composition이 항상 동작함을 가정합니다. 하지만, 항상 성립하는 경우는 아닙니다. 위의 경우도 마찬가지입니다. 눈썰미가 좋으신 분은, try_mapping(F& f, S s, S& r);에서 $f$를 레퍼런스 형태로 받으셨다는 것을 확인할 수 있으실 겁니다. (심지어, 섵불리 수정해서도 안됩니다.) 이는, 함수를 composition을 사용하기 용이한 형태로 만들어주기 위해서입니다.

 

어떤 구간을 나타내는 $s = s_l \circ \cdots \circ s_r$에 대해, 우리가 $f(s)$의 계산에 성공했다고 합시다. 이는 $f(s_l) \circ \cdots \circ f(s_r)$의 계산에 성공했다는 의미입니다. 즉, $f(s_i) = f'(s_i)$인 어떤 $f'$이 존재해서, $f'$은 $\circ$를 보존하고, 항상 계산이 가능했기 때문입니다. 처음 본 문제처럼 $f$가 바뀌지 않을 수도 있습니다. 하지만 이 경우는 $f$가 바뀌는 경우입니다.

 

이 문제에서 $S$에 저장해야 하는 정보는 최솟값, 최댓값, 합, 크기입니다. 최솟값과 최댓값이 같은 경우에만 함수의 적용한다고 하면 계산이 실패하는 경우가 너무 많이 발생할 수 있습니다. $[t^2-1, t^2]$에서 $\lfloor \sqrt{x} \rfloor$를 적용하면 $[t-1, t]$가 되고 여기서 $t^2-t$를 더하면 $[t^2-1, t^2]$으로 돌아옵니다. 그래서 우리는 함수의 적용이 성공하는 경우를 $v = x - \lfloor \sqrt{x} \rfloor$가 모든 $mn$이상 $mx$이하의 수에 대해서 같은 경우 성공했다고 합니다. 이 경우 최솟값과 최댓값의 갱신은 물론, 합의 갱신도 전부다 같은 수인 $v$만큼 빼지기 때문에 $f(s)$를 계산할 수 있습니다. 이 구간에 속하는 수에 대해서는 $\lfloor \sqrt{x} \rfloor = x - v$라는 사실을 알 수 있습니다. 즉, 이 구간에 $f(x) = \sqrt{x}$를 적용하는 대신 $f'(x) = x-v$를 적용해도 상관 없습니다.

 

이를 이용하면, 세그먼트 트리에 lazy하게 남아있는 함수는 모두 $f(x) = x + a$꼴이라는 것을 알 수 있으며, $F$를 $f(x) = x+a$ 혹은 $f(x) = \sqrt{x+a}$꼴로 한정해도 문제가 없습니다.

 

위에서 사용한 구현에서는, 레퍼런스 형태인 $f$를 $f'$으로 바꿔주면 됩니다. 이를 이용해 구현한 수열과 쿼리 28의 코드입니다.

 

#include <bits/stdc++.h>
#include <atcoder/segbeats>
using namespace std;

struct S {
	long long mn, mx, sum, size;
	S() : mn(LONG_MAX), mx(0), sum(0), size(0) {}
	S(int v) : mn(v), mx(v), sum(v), size(1) {}
};
S op(S a, S b) {
	a.mn = min(a.mn, b.mn);
	a.mx = max(a.mx, b.mx);
	a.sum += b.sum, a.size += b.size;
	return a;
}
S e() { return S(); }
using F = pair<bool, long>;
bool try_mapping(F &f, S s, S &r) {
	if (s.size == 0) r = s, return true;
	long long delta = f.second;
	if (f.first) {
		long mn = s.mn + delta, mx = s.mx + delta;
		long dn = (long long)sqrt(mn) - mn, dx = (long long)sqrt(mx) - mx;
		if (dn != dx) return false;
		delta += dn;
	}
	s.mn += delta, s.mx += delta, s.sum += delta * s.size;
	f = {false, delta}, r = s;
	return true;
}
F composition(F f, F g) {
	assert(!g.first);
	return {f.first, f.second + g.second};
}
F id() { return {false, 0L}; }

using Seg = atcoder::segbeats<S, op, e, F, try_mapping, composition, id>;

int main() {
	ios::sync_with_stdio(false); cin.tie(0);
	int N; cin >> N;
	vector<S> V;
	for (int i = 0; i < N; i++) {
		int t; cin >> t;
		V.emplace_back(t);
	}
	Seg seg(V);

	int M; cin >> M;
	for (int i = 0; i < M; i++) {
		int t, l, r;
		cin >> t >> l >> r;
		if (t == 1) {
			int x;
			cin >> x;
			seg.apply(l - 1, r, {false, x});
		} else if (t == 2) seg.apply(l - 1, r, {true, 0});
		else cout << seg.prod(l - 1, r).sum << '\n';
	}
}

문제: ADD, DIV, MAX

그냥 레이지 세그먼트 트리 문제입니다. 함수를 $\left \lfloor \frac{x+a}{b} \right \rfloor$ 형태로 들고 다닙시다.

$\left\lfloor \frac{\left\lfloor x / a \right\rfloor}{b} \right\rfloor = \left\lfloor \frac{x}{ab} \right\rfloor$이기 때문에, 합성이 잘 됩니다.

분자인 $b$가 너무 커지면, 답이 $x$혹은 $x+1$꼴이라는 것을 이용해서 $a$와 $b$를 새롭게 정해줄 수 있습니다. 오버플로우에 유의하면서 구현합시다.

아래 구현은 오버플로우 때문에, 함수를 $\left \lfloor \frac{x+a}{b} \right \rfloor + c$ 형태로 들고다니면서 구현했습니다. $(0 \le a < b)$

using S = int;
S op(S a, S b) { return max(a, b); }
S e() { return 0; }

using F = tuple<int, int, int>;
S mapping(F f, S s) {
	auto [a, b, c] = f;
	return (s + a) / b + c;
}
F composition(F f, F g) {
	auto [fa, fb, fc] = f;
	auto [ga, gb, gc] = g;
	int u = (gc + fa) / fb, d = (gc + fa) % fb;
	long long na = ga + (long long)d * gb, nb = (long long)gb * fb, nc = u + fc;

	const int MAXA = 300'000'001;
	if (nb > MAXA) na = max(0L, MAXA - (nb - na)), nb = MAXA;

	return {na, nb, nc};
}
F id() { return {0L, 1L, 0L}; }