947 Most Stones Removed with Same Row or Column

Link: Most Stones Removed with Same Row or Column

Thought

When I saw the question, I realized its a Union Find question - For any union, I can keep only one stone form it. Then it convert to a count how many unique union finds exist question. 

So, only need to count all unique union find, and use all nodes to minus it. 

Code

class Solution:
    def removeStones(self, stones):
        """
        :type stones: List[List[int]]
        :rtype: int
        """
        N = len(stones)
        nodes = dict()
        cols = collections.defaultdict(set)
        rows = collections.defaultdict(set)
        counter = set()
        for x, y in stones: # pre process
            nodes[(x,y)] = (x,y)
            cols[x].add(y)
            rows[y].add(x)
        def find(x,y):
            if nodes[(x,y)]!= (x,y):
                nx,ny = find(*nodes[(x,y)])
                nodes[(x,y)] = (nx,ny)
            return nodes[(x,y)]
        def union(x,y,x2,y2):
            n_x , n_y = find(x,y)
            nn_x, nn_y = find(x2,y2)
            nodes[(nn_x,nn_y)] = (n_x, n_y)
        for x, y in stones:
            for y2 in cols[x]:
                union(x,y,x,y2)
            for x2 in rows[y]:
                union(x,y, x2,y)
        maxx = 0
        for x,y in stones:
            r_x, r_y = find(x,y)
            counter.add((r_x,r_y))
        return N - len(counter) 

Runtime

Time: O(n^2)

Space: O(n)