924-MinimizeMalwareSpread.go

package main

// 924. Minimize Malware Spread
// You are given a network of n nodes represented as an n x n adjacency matrix graph, 
// where the ith node is directly connected to the jth node if graph[i][j] == 1.

// Some nodes initial are initially infected by malware. Whenever two nodes are directly connected, 
// and at least one of those two nodes is infected by malware, both nodes will be infected by malware. 
// This spread of malware will continue until no more nodes can be infected in this manner.

// Suppose M(initial) is the final number of nodes infected with malware in the entire network after the spread of malware stops. 
// We will remove exactly one node from initial.

// Return the node that, if removed, would minimize M(initial). 
// If multiple nodes could be removed to minimize M(initial), return such a node with the smallest index.

// Note that if a node was removed from the initial list of infected nodes, 
// it might still be infected later due to the malware spread.

// Example 1:
// Input: graph = [[1,1,0],[1,1,0],[0,0,1]], initial = [0,1]
// Output: 0

// Example 2:
// Input: graph = [[1,0,0],[0,1,0],[0,0,1]], initial = [0,2]
// Output: 0

// Example 3:
// Input: graph = [[1,1,1],[1,1,1],[1,1,1]], initial = [1,2]
// Output: 1
 
// Constraints:
//     n == graph.length
//     n == graph[i].length
//     2 <= n <= 300
//     graph[i][j] is 0 or 1.
//     graph[i][j] == graph[j][i]
//     graph[i][i] == 1
//     1 <= initial.length <= n
//     0 <= initial[i] <= n - 1
//     All the integers in initial are unique.

import "fmt"

// Union Find + DFS 
func minMalwareSpread(graph [][]int, initial []int) int {
    n := len(graph)
    parent, size := make([]int, n), make([]int, n)
    for i, _ := range graph {
        parent[i] = i
        size[i] = 1
    }
    //time: O(log N)
    var find func(u int) int
    find = func (u int) int {
        if parent[u] == u {
            return u
        }
        //traversing the path from given node all up to its rep
        leader := find(parent[u])
        //path compression (wuthout it will O(N))
        parent[u] = leader
        return leader
    }
    // union by size -> O(1)
    union := func (u, v int) {
        repU, repV := find(u), find(v)
        if repU != repV {
            if size[repU] > size[repV] {
                parent[repV] = repU // make bigger set the rep of smaller set to minimize height of tree
                size[repU] += size[repV]
            } else {
                parent[repU] = repV
                size[repV] += size[repU]
            }
        }
    }
    min := func (x, y int) int { if x < y { return x; }; return y; }
    for j, _ := range graph {
        for k, _ := range graph[0] {
            if graph[j][k] == 1 {
                union(j, k)
            }
        }
    }
    // identify connected component with 1 infected node
    infectedCountMap := make(map[int]int)
    mn := 1 << 32 - 1
    for i, _ := range initial {
        mn = min(mn, initial[i])
        infectedParent := find(initial[i])
        if val, ok := infectedCountMap[infectedParent]; ok {
            infectedCountMap[infectedParent] = val + 1
        } else {
            infectedCountMap[infectedParent] = 1
        }
    } 
    // choose the biggest connected component 
    res, resSize := mn, -1
    for i, _ := range initial {
        infectedParent := find(initial[i])
        count := infectedCountMap[infectedParent]
        if count == 1 {
            if size[infectedParent] > resSize {
                res = initial[i]
                resSize = size[infectedParent]
            } else if size[infectedParent] == resSize {
                res = min(res, initial[i])
            }
        }
    }
    return res
}

func minMalwareSpread1(graph [][]int, initial []int) int {
    n := len(graph)
    ids, idToSize, id := make([]int, n), make(map[int]int), 0
    for i := range ids {
        if ids[i] == 0 {
            id++
            ids[i] = id
            size := 1
            q := []int{i}
            for len(q) > 0 {
                u := q[0]
                q = q[1:]
                for v := range graph[u] {
                    if ids[v] == 0 && graph[u][v] == 1 {
                        size++
                        q = append(q, v)
                        ids[v] = id
                    }
                }
            }
            idToSize[id] = size
        }
    }
    idToInitials := make(map[int]int)
    for _, u := range initial {
        idToInitials[ids[u]]++
    }
    res := n + 1
    resRemoved := 0
    for _, u := range initial {
        removed := 0
        if idToInitials[ids[u]] == 1 {
            removed = idToSize[ids[u]]
        }
        if removed > resRemoved || (removed == resRemoved && u < res) {
            res, resRemoved = u, removed
        }
    }
    return res
}

func main() {
    fmt.Println(minMalwareSpread([][]int{{1,1,0},{1,1,0},{0,0,1}},[]int{0,1})) // 0
    fmt.Println(minMalwareSpread([][]int{{1,0,0},{0,1,0},{0,0,1}},[]int{0,2})) // 0
    fmt.Println(minMalwareSpread([][]int{{1,1,1},{1,1,1},{1,1,1}},[]int{1,2})) // 1

    fmt.Println(minMalwareSpread1([][]int{{1,1,0},{1,1,0},{0,0,1}},[]int{0,1})) // 0
    fmt.Println(minMalwareSpread1([][]int{{1,0,0},{0,1,0},{0,0,1}},[]int{0,2})) // 0
    fmt.Println(minMalwareSpread1([][]int{{1,1,1},{1,1,1},{1,1,1}},[]int{1,2})) // 1
}