[백준 1321] 군인 C++

1321번: 군인

세그먼트 트리

---

 

10개월 만에 다시 풀어보는 세그먼트 트리문제이다.

 

우선 a ~ b 구간의 부대의 인원의 합을 하나의 노드로 잡고 세그먼트 트리를 만들었다.

 

특정 부대의 증원이나 감원을 할 때는 업데이트 쿼리를 만들어서 수행했다.

ex) 2번 부대의 인원이 늘어났다면 1-5번, 1-3번, 2번 노드의 인원이 모두 같은 수 만큼 증가해야 한다.

 

마지막으로 군인을 배치하는 arrangement쿼리를 만들었다.

군인의 번호가 왼쪽 자식 노드의 최대 인원보다 작다면 왼쪽 트리로 이동한다.

군인의 번호가 왼쪽 자식 노드의 최대 인원보다 크다면 해당 최대 인원만큼 군인의 번호를 빼 준 다음 오른쪽 트리로 이동한다.

리프노드에 도달하면 현재 부대의 번호를 출력한다.

 

 

세그먼트 트리

 

 

오랜만에 세그먼트 트리를 봤는데 예전에 너무 힘들게 공부했어서 그런지 처음에 겁을 많이 먹었다.

그래도 공부한 기억이 남아있어서 바로바로 이해가 됐고 어렵지 않게 풀렸다

과거의 나에게 감사하다

 

 

#include <iostream>
#include <vector>
#include <queue>
#include <algorithm>
#include <cstring>
#include <cmath>
#define fasti ios_base::sync_with_stdio(false); cin.tie(0);
#define fastio ios_base::sync_with_stdio(false); cin.tie(0); cout.tie(0);
#define INF 1e9+7
#define pii pair<int, int>
#define MAX 500001
 
typedef long long ll;
// typedef pair<int, int> pii;
 
using namespace std;
 
 
int N, M;
ll arr[MAX];
ll tree[MAX * 4];
 
void input(){
    cin >> N;
    for(int i = 1; i <= N; i++){
        cin >> arr[i];
    }
}
 
ll make_tree(int startnode, int endnode, int nownode){
    if(startnode == endnode) return tree[nownode] = arr[startnode]; 
    int mid = (startnode + endnode) / 2;
    return tree[nownode] = make_tree(startnode, mid, nownode*2) + \
                           make_tree(mid+1, endnode, nownode*2+1);
}
 
void update(int startnode, int endnode, int targetnode, int nownode, ll val){
    if(targetnode < startnode || targetnode > endnode) return;
    tree[nownode] += val;
    if(startnode == endnode) return;
    int mid = (startnode + endnode) / 2;
    
    update(startnode, mid, targetnode, nownode*2, val);
    update(mid+1, endnode, targetnode, nownode*2+1, val);
}
 
void arrangement(int startnode, int endnode, int nownode, ll soldiernum){
    if(startnode == endnode){
        cout << startnode << "\n";
        return;
    }
    ll left_barracks = tree[nownode*2];
    int mid = (startnode + endnode) / 2;
    
    if(soldiernum <= left_barracks){
        arrangement(startnode, mid, nownode*2, soldiernum);
    }
    else{
        arrangement(mid+1, endnode, nownode*2+1, soldiernum - left_barracks);
    }
}
 
void solve(){
    make_tree(1, N, 1);
    cin >> M;
    int s, i;
    ll v;
    for(int k = 0; k < M; k++){
        cin >> s;
        switch (s){
        case 1:
            cin >> i >> v;
            update(1, N, i, 1, v);
            break;
        case 2:
            cin >> i;
            arrangement(1, N, 1, i);
            
        default: break;
        }
    }
}
 
int main(){
    fastio
    input();
    solve();
    
    return 0;
}