[1321] 군인 - Java[자바]

 

 

| 문제 링크

 

 

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

1321번: 군인

 

 

 

 

 

---

 

 

 

 

 

| 문제

 

 

1. N개의 부대에 인원수가 부여됩니다.
2. 1번 부대부터 1번 군인이 들어가며, 1번 부대가 다 차면 2번 부대에 그 다음 군번의 군인들이 들어갑니다.
3. 1번(자대배치) 과 2번(특정 군번을 가진 군인이 속한 부대를 조회) M번의 명령이 있습니다.
4. 2번 명령이 주어질 때마다 해당 군인이 속한 부대를 출력하시오.

 

 

 

 

---

 

 

 

 

| 풀이

 

우선 이 문제를 풀기 위해서는 이분 탐색과 세그먼트 트리를 이해하고 있어야 합니다. 

 

 

1. 세그먼트 트리는 완전 이진 트리의 모습으로 N개의 배열의 값을 세그먼트 트리로 구현하기 위해서는 N≤2^k를 만족하는 최소한의 k를 구해 세그먼트 트리의 사이즈 2^(k+1)를 구해줍니다.
   N*4를 입력해도 됩니다.
2. 0번 노드는 사용하지 않으며, 루트가 1번 노드를 의미합니다. 또한 해당 세그먼트 트리의 각 노드는 자식 노드들의 합을 의미합니다.
3. 즉 루트(1번 노드)에는 1번 부대부터 N번 부대 까지의 부대원의 합이 들어가 있으며,
   루트의 왼쪽 자식(2번 노드)에는 1번 부대부터 mid((1+N)/2)번 부대까지
   오른쪽 자식(3번 노드)에는 mid+1번 부대부터 N번 부대까지의 부대원의 합이 들어가 있습니다.
4. 3번을 규칙화 하면 i번째 노드는 i*2번째 노드의 값 + i*2+1번째 노드의 값으로 구성되며, 더 이상 부대가 나누어 지지 않는 노드(세그먼트 트리의 가장 밑부분 부터)들은 입력받은 배열의 값이 들어가게 됩니다.
5. 명령으로 1이 들어온 경우 루트부터 해당 부대가 속해있는 부대까지 노드들을 거쳐가며, 노드들의 값을 변경해 줍니다.
6. 명령으로 2가 들어온 경우 루트부터 탐색하며, 다음 두 가지 경우로 재귀 함수를 동작해줍니다.
   1. 왼쪽 자식 노드의 값(해당된 부대들의 군인 수)보다 찾으려는 군번이 작거나 같은 경우
       왼쪽 자식 노드에 해당 군인이 속해 있는 것으로 왼쪽 노드로 이동(mid를 기준으로 왼쪽 부대로 한정됨)하여 재탐색
   2. 왼쪽 자식 노드의 값보다 찾으려는 군번이 큰 경우
       왼쪽 자식 노드에 속한 부대에는 해당 군인을 찾을 수 없으므로 오른쪽 노드로 이동(mid를 기준으로 오른쪽 부대로      한정됨, 왼쪽 자식 노드의 값 만큼 군번을 빼줍니다)하여 재탐색
7. 6번 재귀의 기저 부분으로 하나의 부대만 남은 경우 해당 부대를 출력해 줍니다.

 

 

세그먼트 트리를 글로 이해하시기에는 어려운 부분이 있을거라 예상되며, 주석을 상세하게 작성하였으니 하위 코드가 도움이 되었으면 좋겠습니다.

 

 

 

 

 

---

 

 

 

 

| 코드

 

 

import java.io.*;
import java.util.StringTokenizer;

public class Main {
    
    public static int[] arr;
    public static long[] tree;

    public static void main(String[] args) throws IOException {
        
        BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
        BufferedWriter bw = new BufferedWriter(new OutputStreamWriter(System.out));
        StringBuilder sb = new StringBuilder();
        
        int N = Integer.parseInt(br.readLine());
        StringTokenizer st = new StringTokenizer(br.readLine());
        
        //부대별 병사의 수, 0번 부대는 dafault 값인 0으로 유지한다
        arr = new int[N+1];
        for(int i = 1; i<=N; i++)
            arr[i] = Integer.parseInt(st.nextToken());
        
        //세그먼트 트리를 만들기 위한 size 계산, 2^k>=N을 만족하는 정수 k의 +1을 한 값으로 트리를 생성
        int k = (int) Math.ceil(Math.log(N)/Math.log(2)) + 1;
        int size = (int) Math.pow(2, k);
        
        tree = new long[size];
        
        //이분 탐색으로 각 부대의 인원 수 값을 세그먼트 트리 하단에 채우며, 상위 노드에는 각 부대의 합을 기재한다
        init(1,N,1);
        
        //M개의 명령을 받아 수행
        int t = Integer.parseInt(br.readLine());
        for(int i = 0; i<t; i++) {
            st = new StringTokenizer(br.readLine());
            int a = Integer.parseInt(st.nextToken());
            int b = Integer.parseInt(st.nextToken());
            
            //명령이 1인 경우 b부대의 인원을 c만큼 수정
            if(a == 1) {
            	int c = Integer.parseInt(st.nextToken());
                update(1,N,b,1,c);
            }
            //명령이 2인 경우 b번째 군인이 어느 부대에 속해있는지 출력
            if(a == 2)
            	sb.append(find(1,N,1,b)).append('\n');
        }
        
        bw.write(sb.toString());
        br.close();
        bw.close();
    }
    
    /*
     * s: 첫 부대
     * e: 마지막 부대
     * node: root가 1번 이며, 완전 이진 트리이므로 왼쪽 자식이 node*2, 오른쪽 자식이 node*2 + 1
     * 부모 노드는 자식 노드들의 합을 나타냄
     * root 노드는 s번부터 e번 부대 까지의 합을 나타내고 있으며, root의 왼쪽 자식은 s번부터 mid번, 오른쪽 자식은 mid+1번부터 e번 부대까지의 합을 나타낸다
     * 이러한 완전 이진 트리와 이분 탐색의 특성을 이용해 부대 별 node를 찾는다
     */
    public static long init(int s, int e, int node) {
        
        if(s==e)
            return tree[node] = arr[s];
        
        int mid = (s+e) >>> 1;
        return tree[node] = init(s,mid,node*2) + init(mid+1,e,node*2+1);
    }
    
    
    /*
     * s,e,node는 위와 같음
     * idx: 인원수를 변경하고자하는 부대
     * value: 변경하고자 하는 인원 수
     * idx에 해당하는 부대의 인원을 value만큼 더하고 뺀다
     */
    public static void update(int s, int e, int idx, int node, long value) {
        if(idx<s || idx>e)
            return;
        
        tree[node] += value;
        
        if(s==e)
            return;
        
        int mid = (s+e) >>> 1;
        
        update(s,mid,idx,node*2,value);
        update(mid+1,e,idx,node*2+1,value);
    }
    
    public static int find(int s, int e, int node, long value) {
    	if(s==e)
    		return s;
    	int mid = (s+e) >>> 1;
    	//찾으려는 군인의 번호보다 왼쪽 자식 노드의 값이 크다면, 왼쪽 자식 노드 아래의 부대 안에 군인이 있어 왼쪽 자식 노드로 탐색
    	if(tree[node*2]>=value)
    		return find(s,mid,node*2,value);
    	//찾으려는 군인의 번호가 왼쪽 자식의 노드의 값보다 크다면, 오른쪽 자식 노드 아래의 부대 안에 군인이 있어 오른쪽 자식 노드로 탐색
    	//이때 왼쪽 자식 노드의 값 만큼 군인의 번호를 제외하고 오른쪽 자식 노드에서 탐색
    	else
    		return find(mid+1,e,node*2+1,value-tree[node*2]);
    }

}