[Algorithm] BOJ 13975

BOJ 13975번 파일 합치기 3

Intro

Min Heap을 구현하여 풀이하였습니다.
Swift로 Heap을 직접 구현해본 것은 처음이라 틀린 부분이 있다면 조언 부탁드립니다. Heap 구현 시 해당 블로그를 많이 참고하였습니다.

Solution

1. 입력 받은 리스트를 Min Heap 구조로 변경시킨다.
2. 최솟값을 2번 pop하여 더한다.
3. 더한 값을 Min Heap에 다시 push한다.
4. Min Heap에 남은 요소의 개수가 1개가 될 때까지 2~3번 과정을 반복한다.
5. 남은 요소가 1개면 3번에서 push한 값들의 합을 출력한다.

Code

Swift

import Foundation

let t = Int(readLine()!)!
for _ in 0..<t {
    let _ = Int(readLine()!)!
    let chs = readLine()!.split(separator:" ").map(){Int(String($0))!}
    let heap = MinHeap(elements: chs)
    print(minCost(heap))
}

func minCost(_ heap: MinHeap) -> Int {
    var heap = heap
    var cost = 0
    while heap.count > 1 {
        let ch1 = heap.pop()
        let ch2 = heap.pop()
        heap.insert(ch1+ch2)
        cost += ch1+ch2
    }
    return cost
}

struct MinHeap {
    private var elements: [Int] = []
    init(elements: [Int]) {
        for e in elements { self.insert(e) }
    }
    var count: Int {
        return self.elements.count
    }
    private var lastIndex: Int {
        return self.count - 1
    }
    private func rightChild(of index: Int) -> Int? {
        guard index*2+2 < self.count else { return nil }
        return index*2+2
    }
    private func leftChild(of index: Int) -> Int? {
        guard index*2+1 < self.count else { return nil }
        return index*2+1
    }
    private func parent(of index: Int) -> Int? {
        guard index > 0 else { return nil }
        return (index-1)/2
    }
    private mutating func swap(_ index1: Int, _ index2: Int) {
        let temp = self.elements[index1]
        self.elements[index1] = self.elements[index2]
        self.elements[index2] = temp
    }
    func node(of index: Int) -> Int {
        return self.elements[index]
    }
    mutating func insert(_ node: Int) {
        self.elements.append(node)
        var child = self.lastIndex
        while let parents = parent(of: child),
            self.elements[parents] > node {
            swap(parents, child)
            child = parents
        }
    }
    mutating func delete() {
        swap(self.lastIndex, 0)
        self.elements.remove(at: self.lastIndex)
        var current = 0
        var lChild = leftChild(of: current)
        var rChild = rightChild(of: current)
        while current < self.count, 
        self.elements[current] > self.elements[lChild ?? current] || self.elements[current] > self.elements[rChild ?? current] {
            if let lChild = lChild,
            let rChild = rChild {
                if self.elements[lChild] < self.elements[rChild] {
                    swap(current, lChild)
                    current = lChild
                } else {
                    swap(current, rChild)
                    current = rChild
                }
            } else if let lChild = lChild {
                swap(current, lChild)
                current = lChild
            } else if let rChild = rChild {
                swap(current, rChild)
                current = rChild
            } else { break }
            lChild = leftChild(of: current)
            rChild = rightChild(of: current)
        }
    }
    mutating func pop() -> Int {
        let removeElement = self.elements[0]
        self.delete()
        return removeElement
    }
}

Python

import sys, heapq
input = sys.stdin.readline

def calculate(arr) :
    num = 0
    while len(arr) > 1 :
        ch1 = heapq.heappop(arr)
        ch2 = heapq.heappop(arr)
        num += ch1+ch2
        heapq.heappush(arr, ch1+ch2)
    return num
    
for _ in range(int(input().strip())) :
    k = int(input().strip())
    chs = list(map(int, input().strip().split()))
    heapq.heapify(chs)
    print(calculate(chs))