[백준 1321번] 군인 (C++)

문제링크 : https://www.acmicpc.net/problem/1321

1321번: 군인

#include <bits/stdc++.h>
using namespace std;
typedef long long ll;
typedef pair<int, int> pii;
const int MAX = 987654321;

int N, M;
int arr[500001];
ll Tree[500001 * 4];

void update(int start, int end, int node, int idx, ll diff)
{
    if(idx < start || idx > end) return;  //겹치지 않으므로 종료
    Tree[node] = Tree[node] + diff; //변한만큼 갱신

    if(start!=end)  //리프노드가 아니라면 자식도 갱신해줌
    {
        int mid = (start+end)/2;
        update(start, mid, node*2, idx, diff);    //왼쪽 탐색
        update(mid+1, end, node*2+1, idx, diff);  //오른쪽 탐색
    }
}

void find(int start, int end, int node, ll val)  
{
    if(start == end) 
    {
        cout << start << "\n";
        return;
    }
    int mid = (start + end)/2;
    if(val <= Tree[node * 2]) return find(start, mid, node * 2, val); //왼쪽 자식노드 보다 작다면 왼쪽 탐색
    else return find(mid + 1, end, node*2 + 1, val - Tree[node*2]); //왼쪽 자식노드 보다 크다면 오른쪽 탐색 (해당 인원 만큼 빼줌)
}

int main()
{
    cin >> N;
    for(int i=1; i<=N; i++)
    {
        cin >> arr[i];
        update(1, N, 1, i, arr[i]); //트리 만들기
    }
    
    cin >> M;
    int command, idx;
    ll value;
    for(int i=0; i<M; i++)
    {
        cin >> command;
        if(command == 1) //증감 확인
        {
            cin >> idx >> value;
            update(1, N, 1, idx, value); //증감된 수치만큼 업데이트
        }
        else if(command = 2) //배치 확인
        {
            for(int i=1; i<=N; i++)
            {
                cout << Tree[i] << " ";
            }
            cin >> idx;
            find(1, N, 1, idx); //배치된 곳 찾기
        }
    }

    return 0;
}

세그먼트 트리를 활용한 문제다.

update 를 통해서 초기값과 증감된 수치 업데이트를 해주었다.

 

이후에 배치된 곳을 확인해야 하는데 이 부분은 검색의 도움을 빌렸다.

확인해야 할 곳의 값이 왼쪽 자식 노드 (왼쪽의 최대 인원) 값 보다 작다면 왼쪽을 탐색하고, 아니라면 오른쪽을 탐색한다.

오른쪽을 탐색할 경우에는 왼쪽 최대 값 만큼 빼줘서 남은 값을 탐색한다.

 

처음 할당 값 1, 4, 3, 5, 2
tree[1] = arr[1] ~ arr[5]의 합 : 15
tree[2] = arr[1] ~ arr[3]의 합 : 8
tree[3] = arr[4] ~ arr[5]의 합 : 7
tree[4] = arr[1] ~ arr[2]의 합 : 5
tree[5] = arr[3] : 3

update 
2 -3 / 4 2

tree[1] = 14
tree[2] = 5
tree[3] = 9
tree[4] = 2
tree[5] = 3

찾을 군인 : 5번째

시작 find(1, 5, 1, 5) [start, end, node, val]
(tree[2] = 5) >= 5 -> 왼쪽 탐색 find(1, 3, 4, 5)
(tree[4] = 2) < 5 -> 오른쪽 탐색 (왼쪽 인원만큼 인원 감소 : 5 -> 3) find(3, 3, 9, 3)
start = end로 종료 답은 3