BOJ1321 군인(Java)

---

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

태그 : 세그먼트 트리

---

사고 과정

세그먼트 트리 복습 겸 즐찾 목록에 있던 문제를 하나 집었다.

make_tree나 update는 그냥 일반 세그트리와 동일하게 구현하면 된다.

get 하는 query만 조금 수정해서, 찾아야하는 값을 루트부터 비교한다.

각 노드에서, 현재 노드의 구간 합 값보다 찾는 값이 작거나 같다면 왼쪽으로, 그렇지 않다면 오른쪽으로 이동한다.

왼쪽으로 갈 때는 찾는 값을 그냥 전달하면 되고,

오른쪽으로 갈 때는 (찾는값 - node[left]) 해서 전달해주면 된다.

이후 left==right 될 때 까지 계속 타고 내려가서 정답을 출력하면 된다.

---

코드

import java.io.BufferedReader;
import java.io.IOException;
import java.io.InputStreamReader;
import java.util.StringTokenizer;

public class Main {
	static BufferedReader br = new BufferedReader(new InputStreamReader(System.in));  
	static StringBuilder sb = new StringBuilder();
	static StringTokenizer st;
	
	static int N, M, K;
	static long [] arr = new long [500001];
	static long [] tree = new long [2000004];
	static int ans;
	
	public static void make_tree(int node, int left, int right) {
		if (left == right) {
			tree[node] = arr[left];
			return;
		}
		
		int mid = (left + right) / 2;
		
		make_tree(node * 2, left, mid);
		make_tree(node * 2 + 1, mid + 1, right);
		tree[node] = tree[node * 2] + tree[node * 2 + 1];
	}
	
	public static void update(int node, int left, int right, long idx, long value) {
		if (left == right && idx == left) {
			tree[node] += value;
			return;
		}
		int mid = (left + right) / 2;
		if (idx <= mid) {
			update(node * 2, left, mid, idx, value);
			tree[node] = tree[node * 2] + tree[node * 2 + 1];
		}
		else {
			update(node * 2 + 1, mid + 1, right, idx, value);
			tree[node] = tree[node * 2] + tree[node * 2 + 1];
		}
	}
	
	public static void get(int node, int left, int right, long s, long e, long find) {
		if (right < s || left > e) {
			return;
		}
		
		if (left == right) {
			ans = left;
			return;
		}
		
		int mid = (left + right) / 2;
		
		if (find <= tree[node * 2]) {
			get(node * 2, left, mid, s, e, find);
		}
		else {
			get(node * 2 + 1, mid + 1, right, s, e, find - tree[node * 2]);
		}
		
		
	}
	
	public static void main(String[] args) throws NumberFormatException, IOException {
		st = new StringTokenizer(br.readLine().trim());
		
		N = Integer.parseInt(st.nextToken());
		
		st = new StringTokenizer(br.readLine().trim());

		for(int i = 1; i <= N; ++i) {
			arr[i] = Integer.parseInt(st.nextToken());
		}
		
		make_tree(1, 1, N);
		st = new StringTokenizer(br.readLine().trim());
		M = Integer.parseInt(st.nextToken());
		
		for(int i = 0; i < M; ++i) {
			st = new StringTokenizer(br.readLine().trim());
			long a, b, c;
			a = Long.parseLong(st.nextToken());
		 
			if (a == 1) {
				b = Long.parseLong(st.nextToken());
				c = Long.parseLong(st.nextToken());
				update(1, 1, N, b, c);
			}
			else {
				b = Long.parseLong(st.nextToken());
				get(1, 1, N, 1, N, b);
				sb.append(ans + "\n");
			}
		}
		
		
		System.out.println(sb);
	}
}
 

---

후기

세그트리 복습 겸, 플래 문제도 추가할 겸 좋은 문제였다.

난이도가 그렇게 있지는 않은 것 같았다.