BOJ 1321, 군인

세그먼트트리

BOJ 1321, 군인

이 문제는 세그먼트트리를 이용한 문제입니다.

풀이

이 문제의 핵심은 쿼리부분이라고 생각합니다. 루트노드부터 시작해서 정답노드까지 내려가는 형식입니다.

왼쪽으로 내려갈지 오른쪽으로 내려갈지 판단하는게 중요합니다.

맨 위에 if문 2개는 마지막에 노드값을 출력하기 위한 식이고 밑에 if-else문은 최대한 밑으로 내려가는 식을 구현했습니다.

깔끔하게 풀진 못했지만 참고가 되었으면 좋겠습니다.

#include<stdio.h>
#include<cstring>
#include<algorithm>
#include<string>
#include<iostream>
#include<vector>
#include<cmath>
using namespace std;
typedef long long ll;
int n, h, m;
vector<int> tree;
void update(int idx, int val)
{
	idx += h;
	tree[idx] += val;
	while (idx /= 2)
		tree[idx] = tree[idx * 2 + 1] + tree[idx * 2];
}
int nquery(int node, int val)
{
	if (tree[node] > val && node * 2 >= h * 2)
		return node - h + 1;
	if (tree[node] == val)
	{
		while (node * 2 < h * 2)
		{
			if (tree[node * 2 + 1] != 0)
				node *= 2,node += 1;
			else
				node *= 2;
		}
		return node - h + 1;
	}
	if (val > tree[node * 2])
		return nquery(node * 2 + 1, val - tree[node * 2]);
	else
		return nquery(node * 2, val);
}
int main() {
	scanf("%d", &n);
	h = (1 << (int)log2(n - 1) + 1);
	tree.resize(h * 2, 0);
	for (int i = 0; i < n; i++)
	{
		int value;
		scanf("%d", &value);
		update(i, value);
	}
	scanf("%d", &m);
	for (int i = 1; i <= m; i++)
	{
		int a, b, c;
		scanf("%d", &a);
		if (a == 1)
		{
			scanf("%d %d", &b, &c);
			update(b - 1, c);
		}
		if (a == 2)
		{
			scanf("%d", &b);
			printf("%d\n", nquery(1, b));
		}
	}
	return 0;
}

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