[P5] 백준 1321번 군인 C++ 세그먼트 트리

리뷰

 

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

세그먼트 트리를 활용해 군인을 적절한 부대에 배치시키는 문제

 

 

전역 변수

• N : 배열의 최대 크기를 저장할 상수 변수
• n : 부대의 개수를 저장할 변수
• lst : 각 부대의 병사 수를 저장할 배열
• tree : 세그먼트 트리 정보를 저장할 배열

 

함수

1. build

void build(int node, int s, int e) {
	if (s == e) tree[node] = lst[s];
	else {
		int mid = (s + e) / 2;
		build(node * 2, s, mid);
		build(node * 2 + 1, mid + 1, e);
		tree[node] = tree[node * 2] + tree[node * 2 + 1];
	}
}

 

누적합 세그먼트 트리를 초기화 하기 위한 함수

 

2. update

void update(int node, int s, int e, int idx, int val) {
	if (s == e) tree[node] += 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] = tree[node * 2] + tree[node * 2 + 1];
	}
}

 

세그먼트 트리의 특정 인덱스에 해당하는 노드값을 val만큼 변경하는 함수

 

3. query

int query(int node, int s, int e, int idx) {
	if (s == e) return s;
	int mid = (s + e) / 2;
	if (idx <= tree[node * 2]) return query(node * 2, s, mid, idx);
	return query(node * 2 + 1, mid + 1, e, idx - tree[node * 2]);
}

 

idx번째 군인이 몇번 부대에 배치될지를 찾는 함수

 

문제풀이

1. n값을 입력 받고, n개 부대의 군인 수를 lst배열에 입력 받아준다.
2. build함수를 통해 재귀를 사용하여 세그먼트 트리의 리프노드에 배열의 값을 저장해 준다.
3. 재귀를 빠져나오면서 좌, 우 자식 노드의 합을 상위 노드의 값으로 저장해 준다.
4. m값을 입력 받고, m개의 쿼리를 수행하여 매 쿼리마다 변수 op, idx에 값을 입력받아 준다.
5. op가 1일 경우 변수 val에 추가로 값을 입력 받아주고, update함수를 통해 세그먼트 트리의 idx을 val만큼 변경해준다.
6. op가 2일 경우 query함수를 통해 세그먼트 트리의 idx번째 요소가 저장된 노드의 인덱스를 찾아 출력해 준다.

 

트러블 슈팅

없음

 

 

참고 사항

• 쿼리 함수의 idx가 좌측 자식 노드보다 크다면 idx를 좌측 자식 노드의 값만큼 빼주면서 우측 자식 노드로 이동한다.
• idx가 좌측 자식노드보다 작다면 idx를 그대로 좌측 노드로 이동해 준다.

 

정답 코드

#include<iostream>
using namespace std;

const int N = 500001;
int n, m;
int lst[N];
int tree[N * 4];

void build(int node, int s, int e) {
	if (s == e) tree[node] = lst[s];
	else {
		int mid = (s + e) / 2;
		build(node * 2, s, mid);
		build(node * 2 + 1, mid + 1, e);
		tree[node] = tree[node * 2] + tree[node * 2 + 1];
	}
}

void update(int node, int s, int e, int idx, int val) {
	if (s == e) tree[node] += 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] = tree[node * 2] + tree[node * 2 + 1];
	}
}

int query(int node, int s, int e, int idx) {
	if (s == e) return s;
	int mid = (s + e) / 2;
	if (idx <= tree[node * 2]) return query(node * 2, s, mid, idx);
	return query(node * 2 + 1, mid + 1, e, idx - tree[node * 2]);
}

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

	cin >> n;
	for (int i = 1; i <= n; ++i) cin >> lst[i];
	build(1, 1, n);

	cin >> m;
	while (m--) {
		int op, idx; cin >> op >> idx;
		if (op == 1) {
			int val; cin >> val;
			update(1, 1, n, idx, val);
		}
		else {
			cout << query(1, 1, n, idx) << "\n";
		}
	}
}