글 읽기 - [Swift] 33%에서 시간초과..! Prim 알고리즘사용

yabby1997 3년 전 0

Prim 알고리즘으로 MST를 총 두번, 파란간선을 최소로 선택하는 경우, 최대로 선택하는경우를 구해준 뒤

주어진 파란간선의 개수가 그 사이에 위치하는지를 확인하였습니다.

그런데 33%에서 시간초과가 발생하네요..

처음엔 무식하게 모든 경우를 일일히 확인하도록 했는데 이때도 33%에서 시간초과..

Kruskal을 이용했을 때에도 33%에서 시간초과가 발생해서

마지막 희망을 갖고 Prim으로 구현해보았는데도 33%에서 시간초과가 발생하니.... 절망이네요 ㅜㅜ

고수님들의 도움 부탁드립니다 ..!

import Foundation

struct EdgeData: Comparable {
    let to: Int
    let cost: Int
    
    static func < (lhs: EdgeData, rhs: EdgeData) -> Bool {
        lhs.cost < rhs.cost
    }
}

class PriorityQueue<T> {
    var heap: [T]
    let priority: (T, T) -> Bool
    
    init(heap: [T], priority: @escaping (T, T) -> Bool) {
        self.heap = heap
        self.priority = priority
        buildHeap()
    }
    
    func isEmpty() -> Bool {
        return heap.isEmpty
    }
    
    func size() -> Int {
        return heap.count
    }
    
    func leftChildIndex(of index: Int) -> Int {
        return index * 2 + 1
    }
    
    func rightChildIndex(of index: Int) -> Int {
        return index * 2 + 2
    }
    
    func parentIndex(of index: Int) -> Int {
        return (index - 1) / 2
    }
    
    func enqueue(_ element: T) {
        var currentIndex = heap.count
        heap.append(element)
        
        while currentIndex != 0 {
            let parent = parentIndex(of: currentIndex)
            if priority(heap[currentIndex], heap[parent]) {
                heap.swapAt(currentIndex, parent)
                currentIndex = parent
            } else {
                break
            }
        }
    }
    
    func dequeue() -> T?{
        if isEmpty() {
            return nil
        } else {
            heap.swapAt(0, heap.count - 1)
            let target = heap.removeLast()
            heapify(at: 0)
            return target
        }
    }
    
    func heapify(at index: Int) {
        var targetIndex = rightChildIndex(of: index)
        let leftIndex = leftChildIndex(of: index)
        
        if targetIndex < heap.count {
            if priority(heap[leftIndex], heap[targetIndex]) {
                targetIndex = leftIndex
            }
        } else if leftIndex < heap.count {
            targetIndex = leftIndex
        } else {
            return
        }
        
        if priority(heap[targetIndex], heap[index]) {
            heap.swapAt(targetIndex, index)
            heapify(at: targetIndex)
        } else {
            return
        }
    }
    
    func buildHeap() {
        for eachParent in (0..<(heap.count / 2)).reversed() {
            heapify(at: eachParent)
        }
    }
}

var adjacencyList = [[EdgeData]]()
var numOfVertices = 0
var numOfEdges = 0
var givenNumOfBlueEdges = 0

while true {
    let input = readLine()!.split(separator: " ").map({Int(String($0))!})
    numOfVertices = input[0]
    numOfEdges = input[1]
    givenNumOfBlueEdges = input[2]
    
    adjacencyList = Array(repeating: [EdgeData](), count: numOfVertices + 1)
    
    if numOfVertices == 0 && numOfVertices == numOfEdges && numOfEdges == givenNumOfBlueEdges {
        break
    }
    
    for _ in 0..<numOfEdges {
        let eachEdgeInfo = readLine()!.split(separator: " ").map({String($0)})
        let colorOfEdge = eachEdgeInfo[0]
        let from = Int(eachEdgeInfo[1])!
        let to = Int(eachEdgeInfo[2])!
        
        if colorOfEdge == "B" {
            adjacencyList[from].append(EdgeData(to: to, cost: 1))
            adjacencyList[to].append(EdgeData(to: from, cost: 1))
        } else {
            adjacencyList[from].append(EdgeData(to: to, cost: 0))
            adjacencyList[to].append(EdgeData(to: from, cost: 0))
        }
    }
    
    let minBlueEdges = prim(priority: <)
    let maxBlueEdges = prim(priority: >)
    
    if minBlueEdges <= givenNumOfBlueEdges && givenNumOfBlueEdges <= maxBlueEdges {
        print(1)
    } else {
        print(0)
    }
}

func prim(priority: @escaping (EdgeData, EdgeData) -> Bool) -> Int {
    let priorityQueue = PriorityQueue<EdgeData>(heap: [], priority: priority)
    var visited = Array(repeating: false, count: numOfVertices + 1)
    var totalCost = 0
    
    visited[1] = true
    for i in 0..<adjacencyList[1].count {
        priorityQueue.enqueue(adjacencyList[1][i])
    }
    
    while !priorityQueue.isEmpty() {
        let currentEdge = priorityQueue.dequeue()!
        let to = currentEdge.to
        let cost = currentEdge.cost
        
        if visited[to] {
            continue
        } else {
            visited[to] = true
            for i in 0..<adjacencyList[to].count {
                priorityQueue.enqueue(adjacencyList[to][i])
            }
            totalCost += cost
        }
    }
    
    for index in 1..<visited.count {
        if !visited[index] {
            return 1000000
        }
    }
    return totalCost
}

댓글을 작성하려면 로그인해야 합니다.

문제

문제

출처

ICPC

[Swift] 33%에서 시간초과..! Prim 알고리즘사용

yabby1997 3년 전 0

Baekjoon Online Judge

채점 현황

문제

유저 대회

고등학교 대회

출처

대학교 대회

도움말