백준 2243번 사탕상자 - 스위프트(Swift) 풀이

 CLASS 6 따기 중인데, 이제 3문제 밖에 안 남았다:)

 

 세그먼트 트리 문제. 구간합 문제라 더 구현이 간단한 펜윅트리로 풀었다. 펜윅트리 구현만 할 줄 알면 단순 적용만 하면 된다. 플래티넘 5 문제 중에는 이렇게 조금 난이도 있는 자료구조/알고리즘을 단순 적용하는 문제들이 꽤 있다.

 

 사탕들은 아래와 같이 관리한다. 

candy[i] = 맛 순위가 i인 사탕 개수

 

 먼저 n번째 사탕을 찾을 때를 보면, n번째로 맛있는 사탕의 맛 순위는 candy[1...k]의 구간합이 최초로 n이상이 되는 k(lower bound)가 된다. 사탕을 넣거나 뺄 때는 candy[i]에 원하는 개수만큼을 더하거나 빼면 된다. 이렇게 계속 변화하는 배열의 구간합을 구해야 하므로, 펜윅트리를 쓸 수 있다.

 

	import Foundation
	class FenwickTree {
	var tree = Array(repeating: 0, count: 1000001)
	func nthCandy(n: Int) -> Int {
	var left = 1, right = 1000000
	while (left < right) {
	let mid = (left + right) / 2
	if n <= sum(to: mid) { right = mid }
	else { left = mid + 1 }
	}
	return left
	}
	func sum(to: Int) -> Int {
	var ret = 0
	var pos = to
	while (pos > 0) {
	ret += tree[pos]
	pos &= (pos - 1)
	}
	return ret
	}
	func add(at: Int, val: Int) {
	var pos = at
	while(pos < tree.count) {
	tree[pos] += val
	pos += (pos & -pos)
	}
	}
	}
	let fwt = FenwickTree()
	var result = ""
	let n = Int(readLine()!)!
	for _ in 0..<n {
	let query = readLine()!.split(separator: " ").map{ Int(String($0))! }
	switch query[0] {
	case 1:
	let rank = query[1]
	let nthCandy = fwt.nthCandy(n: rank)
	result.write("\(nthCandy)\n")
	fwt.add(at: nthCandy, val: -1)
	case 2:
	let flavor = query[1]
	let count = query[2]
	fwt.add(at: flavor, val: count)
	default:
	break
	}
	}
	print(result)

view raw BOJ_2243.swift hosted with ❤ by GitHub