1395번: 스위치 
문제 링크: 1395번: 스위치
개요
세그먼트 트리를 이용해 쿼리를 수행하고 부분합을 출력하는 문제이다. 그러나 원소의 개수와 쿼리의 개수가 최대 100,000으로, 평범한 세그먼트 트리로 풀이하면 시간 초과를 맞이하는 미래가 찾아온다. 이 문제는 ‘느리게 갱신되는 세그먼트 트리’ 알고리즘을 이용할 것이다.
평범한 세그먼트 트리에 대해 잘 모른다면 [백준 14438번] 수열과 쿼리 17 [C++] 포스팅을 읽어보길 바란다.
필자는 세그먼트 트리와 같은 크기의 lazy 배열을 이용할 것이다. (구간합을 저장하는 세그먼트 트리의 원소를 pair<int, int>로 바꿔 하나의 배열에 같이 저장해도 상관없다.) 평범한 세그먼트 트리에서는 매 구간 업데이트 쿼리마다 세그먼트 트리의 리프 노드까지 업데이트했다. 그러나 시간을 절약하기 위해 우린 조금 게을러질 필요가 있다. 업데이트 된 정보를 바로바로 자식 노드에게 물려주는 게 아니라 lazy라는 배열에 저장만 해뒀다가, “필요할 때” 업데이트하는 것이다.
풀이
구간 합 반환 함수
1
2
3
4
5
6
7
8
9
int query(int node, int start, int end, int L, int R) {
applyLazy(node, start, end);
if (end < L || R < start) return 0;
if (L <= start && end <= R) return segTree[node];
int mid = (start + end) / 2;
return query(node * 2, start, mid, L, R) + query(node * 2 + 1, mid + 1, end, L, R);
}
일반적인 세그먼트 트리의 구간 합 함수와 동일하다. 현재 탐색 구간(start~end)이 목표 구간(L~R)에 전혀 포함되지 않으면 함수 탈출, 완전히 포함돠면 해당 세그먼트 트리의 값 반환, 일부만 겹치만 탐색 구간을 반으로 나눠서 재귀를 통해 탐색하는 과정을 거친다.
다른 점이 있다면 맨 처음마다 applyLazy 함수를 호출한다는 것이다.
lazy 적용하는 함수
1
2
3
4
5
6
7
8
9
10
11
12
13
void applyLazy(int node, int start, int end) {
if (lazy[node] == 0) return;
segTree[node] = (end - start + 1) - segTree[node];
if (start != end) {
lazy[node * 2] ^= 1;
lazy[node * 2 + 1] ^= 1;
}
lazy[node] = 0;
return;
}
현재 노드에 저장된 lazy값이 있다면(0이 아니라면) 현재 노드의 값을 업데이트하고, 자식 노드에게 lazy를 물려준다. 이후 중복 작업이 일어나지 않게 현재 노드의 lazy를 0으로 초기화해준다. 이 문제의 쿼리는 구간 반전이기 때문에 목적에 맞게 수행해주었다.
세그먼트 트리 업데이트 함수
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
void updateSegTree(int node, int start, int end, int L, int R) {
applyLazy(node, start, end);
if (L > end || R < start) return;
if (L <= start && end <= R) {
segTree[node] = (end - start + 1) - segTree[node];
if (start != end) {
lazy[node * 2] ^= 1;
lazy[node * 2 + 1] ^= 1;
}
return;
}
int mid = (start + end) / 2;
updateSegTree(node * 2, start, mid, L, R);
updateSegTree(node * 2 + 1, mid + 1, end, L, R);
segTree[node] = segTree[node * 2] + segTree[node * 2 + 1];
return;
}
일반적인 세그먼트 트리 업데이트 함수와 비슷하다. 그러나 구간 합 반환 함수와 마찬가지로 맨 처음마다 applyLazy 함수를 호출해줘야 한다. 또한 탐색 구간이 목표 구간에 환전히 포함됐을 때 현재 노드의 값을 업데이트 것 뿐만 아니라 현재 노드의 자식 노드의 lazy 값을 수정해줘야 한다.
전체 코드
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
#include <iostream>
#pragma warning(disable:4996)
#pragma warning(disable:6031)
using namespace std;
#define MAX_N 1000000
int n, k;
int arr[MAX_N + 1] = { 0 }, segTree[MAX_N * 4] = { 0 }, lazy[MAX_N * 4] = { 0 };
int query(int node, int start, int end, int L, int R);
void applyLazy(int node, int start, int end);
void updateSegTree(int node, int start, int end, int L, int R);
int main() {
scanf("%d %d", &n, &k);
for (int t = 0; t < k; t++) {
int cmd, a, b;
scanf("%d %d %d", &cmd, &a, &b);
a--, b--;
if (cmd == 0) updateSegTree(1, 0, n - 1, a, b);
else if (cmd == 1) printf("%d\n", query(1, 0, n - 1, a, b));
}
return 0;
}
int query(int node, int start, int end, int L, int R) {
applyLazy(node, start, end);
if (end < L || R < start) return 0;
if (L <= start && end <= R) return segTree[node];
int mid = (start + end) / 2;
return query(node * 2, start, mid, L, R) + query(node * 2 + 1, mid + 1, end, L, R);
}
void applyLazy(int node, int start, int end) {
if (lazy[node] == 0) return;
segTree[node] = (end - start + 1) - segTree[node];
if (start != end) {
lazy[node * 2] ^= 1;
lazy[node * 2 + 1] ^= 1;
}
lazy[node] = 0;
return;
}
void updateSegTree(int node, int start, int end, int L, int R) {
applyLazy(node, start, end);
if (L > end || R < start) return;
if (L <= start && end <= R) {
segTree[node] = (end - start + 1) - segTree[node];
if (start != end) {
lazy[node * 2] ^= 1;
lazy[node * 2 + 1] ^= 1;
}
return;
}
int mid = (start + end) / 2;
updateSegTree(node * 2, start, mid, L, R);
updateSegTree(node * 2 + 1, mid + 1, end, L, R);
segTree[node] = segTree[node * 2] + segTree[node * 2 + 1];
return;
}