99% Python union find with optimization
Intuition\n Describe your first thoughts on how to solve this problem. \n\n# Approach\n Describe your approach to solving the problem. \n\n# Complexity\n- Time 
# Intuition\n<!-- Describe your first thoughts on how to solve this problem. -->\n\n# Approach\n<!-- Describe your approach to solving the problem. -->\n\n# Complexity\n- Time complexity:\n<!-- Add your time complexity here, e.g. $$O(n)$$ -->\n\n- Space complexity:\n<!-- Add your space complexity here, e.g. $$O(n)$$ -->\n\n# Code\n```\nclass UnionFind:\n    def __init__(self, n):\n        self.parent = list(range(n))\n        self.size = [1] * n\n\n    def find(self, x):\n        if x != self.parent[x]:\n            self.parent[x] = self.find(self.parent[x])\n        return self.parent[x]\n\n    def union(self, x, y):\n        a, b = self.find(x), self.find(y)\n        if a != b:\n            if self.size[a] < self.size[b]:\n                self.parent[a] = b\n                self.size[b] += self.size[a]\n            else:\n                self.parent[b] = a\n                self.size[a] += self.size[b]\n\nclass Solution:\n    def minMalwareSpread(self, graph, initial):\n        n, ans, max_save = len(graph), min(initial), 0\n        infected = [False] * n\n        for u in initial:\n            infected[u] = True\n        uf = UnionFind(n)\n        \n        # Union all non-infected nodes\n        for i in range(n):\n            if not infected[i]:\n                for j in range(i + 1, n):\n                    if not infected[j] and graph[i][j] == 1:\n                        uf.union(i, j)\n        \n        # Count the influence of each initially infected node\n        component = [set() for _ in range(n)]\n        count = [0] * n\n        for u in initial:\n            component[u] = {uf.find(v) for v in range(n) if not infected[v] and graph[u][v] == 1}\n            for v in component[u]:\n                count[v] += 1\n        \n        # Determine the node removal that saves the most nodes\n        for u in sorted(initial):\n            save = sum(uf.size[v] for v in component[u] if count[v] == 1)\n            if save > max_save or (save == max_save and u < ans):\n                max_save, ans = save, u\n        \n        return ans\n\n\n```
['Python3']