https://www.acmicpc.net/problem/1321
세그먼트트리 문제다. 바로 이전 문제인 [2243] 사탕상자 문제와 동일한 풀이로 접근하면 된다.
풀이
N개의 부대에 속한 군인 수를 저장하는 구간 합 세그먼트트리를 생성한다. 이후 i번째 군인의 부대를 호출할 때, 현재 구간의 군인 수가 i보다 크다면 좌측노드에서 i번째 군인 탐색. 현재 구간의 군인 수가 i보다 적다면 (i - 좌측구간 군인수) 번째 군인을 우측노드에서 탐색한다. 없는 군인을 탐색하는 경우는 주어지지 않으므로 리프노드에 도달할 경우에 해당 노드의 인덱스를 반환하면 된다.
정답 코드
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| import Foundation | |
| let N = Int(readLine()!)! | |
| let arr = readLine()!.split(separator: " ").map{Int($0)!} | |
| var tree = Array(repeating: 0, count: N*4) | |
| func make_tree(from:Int, to:Int, idx:Int) -> Int{ | |
| if from == to{ | |
| tree[idx] = arr[from] | |
| return tree[idx] | |
| } | |
| let mid = (from+to)/2 | |
| let left = make_tree(from: from, to: mid, idx: idx*2) | |
| let right = make_tree(from: mid+1, to: to, idx: idx*2+1) | |
| tree[idx] = left+right | |
| return tree[idx] | |
| } | |
| func update_tree(from:Int, to:Int, idx:Int, target_idx:Int, value:Int){ | |
| if target_idx < from || to < target_idx{ | |
| return | |
| } | |
| tree[idx] += value | |
| if from == to{ | |
| return | |
| } | |
| let mid = (from+to)/2 | |
| update_tree(from: from, to: mid, idx: idx*2, target_idx: target_idx, value: value) | |
| update_tree(from: mid+1, to: to, idx: idx*2+1, target_idx: target_idx, value: value) | |
| } | |
| func soldier_num(from:Int, to:Int, idx:Int, num:Int) -> Int{ | |
| if from == to{ | |
| return from | |
| } | |
| let mid = (from+to)/2 | |
| let left = tree[idx*2] | |
| if left >= num{ | |
| return soldier_num(from: from, to: mid, idx: idx*2, num: num) | |
| }else{ | |
| return soldier_num(from: mid+1, to: to, idx: idx*2+1, num: num-left) | |
| } | |
| } | |
| make_tree(from: 0, to: N-1, idx: 1) | |
| var ans = [String]() | |
| let M = Int(readLine()!)! | |
| for _ in 0..<M{ | |
| let line = readLine()!.split(separator: " ").map{Int($0)!} | |
| let cmd = line[0] | |
| if cmd == 1{ | |
| let i = line[1]-1 | |
| let v = line[2] | |
| update_tree(from: 0, to: N-1, idx: 1, target_idx: i, value: v) | |
| }else{ | |
| let i = line[1] | |
| ans.append("\(soldier_num(from: 0, to: N-1, idx: 1, num: i)+1)") | |
| } | |
| } | |
| print(ans.joined(separator: "\n")) |
'Problem Solving > BOJ' 카테고리의 다른 글
| [1517] 버블 소트 (0) | 2023.03.02 |
|---|---|
| [1572] 중앙값 (0) | 2023.03.01 |
| [2243] 사탕상자 (0) | 2023.02.28 |
| [3745] 오름세 (0) | 2023.02.26 |
| [12837] 가계부 (Hard) (0) | 2023.02.23 |
