[P4] 백준 16221번 모독 C++ 세그먼트 트리

리뷰

 

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

세그먼트 트리 누적합의 범위를 어떻게 해야할지 좀 오래 고민을 하게 되었던 문제였다.

 

 

전역 변수

• N : 배열의 최대 크기를 정의할 상수 변수
• q : 쿼리의 개수를 저장할 변수
• t : 쿼리의 연산 번호를 저장할 변수
• k : 수열에 추가할 원소를 저장할 변수
• Node : 세그먼트 트리 노드 정보를 정의할 구조체
• tree : 세그먼트 트리 구간합과 구간 최소값을 저장할 Node타입 배열

 

함수

1. update

void update(int node, int s, int e, int idx, int val) {
	if (s == e) tree[node].sum += val, tree[node].mn += val;
	else {
		int mid = (s + e) / 2;
		if (idx <= mid) update(node * 2, s, mid, idx, val);
		else update(node * 2 + 1, mid + 1, e, idx, val);
		tree[node].sum = tree[node * 2].sum + tree[node * 2 + 1].sum;
		tree[node].mn = min(tree[node * 2].mn, tree[node * 2 + 1].mn);
	}
}

 

2. get_idx

int get_idx(int node, int s, int e) {
	if (s == e) return s;
	int mid = (s + e) / 2;
	if (tree[node * 2].mn == 0) return get_idx(node * 2, s, mid);
	else if (tree[node * 2 + 1].mn == 0) return get_idx(node * 2 + 1, mid + 1, e);
	return N;
}

 

세그먼트 트리 구간의 리프 노드 값이 0인 가장 작은 인덱스를 구하기 위한 함수

 

3. query

int query(int node, int s, int e, int L, int R) {
	if (R < s || e < L) return 0;
	if (L <= s && e <= R) return tree[node].sum;
	int mid = (s + e) / 2;
	return query(node * 2, s, mid, L, R) + query(node * 2 + 1, mid + 1, e, L, R);
}

 

세그먼트 트리 구간합을 구하기 위한 함수

 

 

문제풀이

1. 변수 q에 값을 입력 받고, q번의 쿼리를 수행한다.
2. 매 쿼리마다 변수 t, k에 값을 입력 받는다.
3. t가 1일 경우 세그먼트 트리의 k번 인덱스 리프노드에 1을 더해주고, 부모 노드를 업데이트한다.
4. t가 2일 경우 세그먼트 트리의 k번 인덱스 리프노드에 -1을 더해주고, 부모 노드를 업데이트한다.
5. 변수 idx에 get_idx함수를 통해 1~N범위에서 리프 노드 값이 0인 가장 작은 인덱스를 저장한다.
6. query함수를 통해 세그먼트 트리의 1~idx번째 노드의 구간합을 구해 출력하고, 개행문자를 삽입한다.

 

트러블 슈팅

없음

 

 

참고 사항

1. (1 ≤ Q ≤ 1,000,000), (T = 1 또는 T = 2, 1 ≤ K ≤ 1,000,000)로 세그먼트 트리로 풀이가 가능하다.
2. 1부터 연속된 구간의 최대 인덱스를 구하기 위해 세그먼트 트리의 노드에 최소값을 추가로 구현하였다.

 

정답 코드

#include<iostream>
using namespace std;

const int N = 1e6 + 1;
int q, t, k;
struct Node {
	int sum, mn;
};
Node tree[N * 4];

void update(int node, int s, int e, int idx, int val) {
	if (s == e) tree[node].sum += val, tree[node].mn += val;
	else {
		int mid = (s + e) / 2;
		if (idx <= mid) update(node * 2, s, mid, idx, val);
		else update(node * 2 + 1, mid + 1, e, idx, val);
		tree[node].sum = tree[node * 2].sum + tree[node * 2 + 1].sum;
		tree[node].mn = min(tree[node * 2].mn, tree[node * 2 + 1].mn);
	}
}

int get_idx(int node, int s, int e) {
	if (s == e) return s;
	int mid = (s + e) / 2;
	if (tree[node * 2].mn == 0) return get_idx(node * 2, s, mid);
	else if (tree[node * 2 + 1].mn == 0) return get_idx(node * 2 + 1, mid + 1, e);
	return N;
}

int query(int node, int s, int e, int L, int R) {
	if (R < s || e < L) return 0;
	if (L <= s && e <= R) return tree[node].sum;
	int mid = (s + e) / 2;
	return query(node * 2, s, mid, L, R) + query(node * 2 + 1, mid + 1, e, L, R);
}

int main() {
	ios::sync_with_stdio(0);
	cin.tie(0);
	cout.tie(0);

	cin >> q;
	while (q--) {
		cin >> t >> k;
		if (t == 1) update(1, 1, N, k, 1);
		else update(1, 1, N, k, -1);

		int idx = get_idx(1, 1, N);
		cout << query(1, 1, N, 1, idx) << "\n";
	}
}