[BOJ / Swift & Python] 13975 - 파일 합치기 3

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))

* 개인 블로그에 작성한 내용을 복사해왔습니다.