99클럽 코테 스터디 21일차 TIL: 힙 응용

문제

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

풀이

힙을 사용하여 풀 수 있는 문제이다.
최대값을 계속해서 뽑아내야 하므로 최대힙을 사용하는 것이 시간복잡도 측에서 유리한 방법이다.

최대힙으로 pop 연산을 하여 입력받은 h보다 작다면 YES 및 횟수 아니라면 NO 및 최대값을 출력해주면 된다.
한가지 고려해야할 점으로는 1/2 연산을 하기도 전에 모든 거인의 키가 입력받은 h보다 작을 때를 고려해주어야한다.

소스코드

	import Foundation
	struct Heap<T: Comparable> {
	private var elements: [T] = []
	private let comparer: (T, T) -> Bool
	init(comparer: @escaping (T,T) -> Bool) {
	self.comparer = comparer
	}
	mutating func insert(element: T) {
	if elements.isEmpty {
	elements.append(element)
	elements.append(element)
	return
	}
	elements.append(element)
	swimUp(index: elements.count - 1)
	}
	mutating private func swimUp(index: Int) {
	var index = index
	while index > 1 && comparer(elements[index], elements[index / 2]) {
	elements.swapAt(index, index / 2)
	index /= 2
	}
	}
	mutating func pop() -> T? {
	if elements.count < 2 { return nil }
	elements.swapAt(1, elements.count - 1)
	let deletedElement = elements.removeLast()
	diveDown(index: 1)
	return deletedElement
	}
	mutating private func diveDown(index: Int) {
	var swapIndex = index
	var isSwap = false
	let leftIndex = index * 2
	let rightIndex = index * 2 + 1
	if leftIndex < elements.endIndex && comparer(elements[leftIndex], elements[swapIndex]) {
	swapIndex = leftIndex
	isSwap = true
	}
	if rightIndex < elements.endIndex && comparer(elements[rightIndex], elements[swapIndex]) {
	swapIndex = rightIndex
	isSwap = true
	}
	if isSwap {
	elements.swapAt(swapIndex, index)
	diveDown(index: swapIndex)
	}
	}
	}
	let input = readLine()!.split { $0 == " " }.map { Int($0)! }
	let n = input[0], h = input[1], t = input[2]
	var heap = Heap<Int>(comparer: >)
	for _ in 0..<n { heap.insert(element: Int(readLine()!)!) }
	var isSuccess = false
	var count = 0
	for _ in 0..<t {
	let element = heap.pop()!
	if element < h {
	heap.insert(element: element)
	break
	}
	count += 1
	heap.insert(element: max(element / 2, 1))
	}
	let top = heap.pop()!
	if top < h {
	print("YES")
	print(count)
	} else {
	print("NO")
	print(top)
	}

view raw 19638_센티와 마법의 뿅망치.swift hosted with ❤ by GitHub

후기

고려해야할 점을 몰랐다면 맞외틀이 나올만한 문제였다..