928. Minimize Malware Spread II

leetcode928.swift

class Solution {
    func minMalwareSpread(_ graph: [[Int]], _ initial: [Int]) -> Int {
        
        var record = [Int:[Int]]()
        
        let initialSet = Set(initial)
        
        for start in initial {
            var visited = initialSet
            
            var queue = [start]
            
            while !queue.isEmpty {
                var nextQueue = [Int]()
                
                for node in queue {
                    let array = graph[node]
                    for i in 0..<array.count where array[i] == 1 && !visited.contains(i) {
                        nextQueue.append(i)
                        visited.insert(i)
                        record[i, default:[Int]()].append(start)
                    }
                }
                
                queue = nextQueue
            }
        }
        
        var result = [Int](repeating:0 ,count: graph.count)
        
        for (key, value) in record where value.count == 1 {
            for node in value {
                result[node] += 1
            }
        }

        let maxValue = result.max() ?? 0
        
        for i in initial.sorted() where result[i] == maxValue {
            return i
        }
        
        return initial[0]
    }
}