<p>You are given an <code>n x n</code> binary matrix <code>grid</code>. You are allowed to change <strong>at most one</strong> <code>0</code> to be <code>1</code>.</p>

<p>Return <em>the size of the largest <strong>island</strong> in</em> <code>grid</code> <em>after applying this operation</em>.</p>

<p>An <strong>island</strong> is a 4-directionally connected group of <code>1</code>s.</p>

<p>&nbsp;</p>
<p><strong class="example">Example 1:</strong></p>

<pre>
<strong>Input:</strong> grid = [[1,0],[0,1]]
<strong>Output:</strong> 3
<strong>Explanation:</strong> Change one 0 to 1 and connect two 1s, then we get an island with area = 3.
</pre>

<p><strong class="example">Example 2:</strong></p>

<pre>
<strong>Input:</strong> grid = [[1,1],[1,0]]
<strong>Output:</strong> 4
<strong>Explanation: </strong>Change the 0 to 1 and make the island bigger, only one island with area = 4.</pre>

<p><strong class="example">Example 3:</strong></p>

<pre>
<strong>Input:</strong> grid = [[1,1],[1,1]]
<strong>Output:</strong> 4
<strong>Explanation:</strong> Can&#39;t change any 0 to 1, only one island with area = 4.
</pre>

<p>&nbsp;</p>
<p><strong>Constraints:</strong></p>

<ul>
	<li><code>n == grid.length</code></li>
	<li><code>n == grid[i].length</code></li>
	<li><code>1 &lt;= n &lt;= 500</code></li>
	<li><code>grid[i][j]</code> is either <code>0</code> or <code>1</code>.</li>
</ul>

from collections import defaultdict
class Solution:
    def largestIsland(self, grid: List[List[int]]) -> int:
        # for each 0 value, change it to 1, check the largest island 
        # use union find
        def find(x):
            y = parent[x][0]
            if y!=x:
                y = find(y)
                parent[x][0] = y
            return y
        def union(x,y):  # x add to y
            pX = find(x)
            pY = find(y)
            if pX == pY:
                return
            parent[pX][0]  = parent[pY][0] # modify ID
            parent[pY][1] += parent[pX][1] # modify Area
        
        def getID(i,j):
            nonlocal n
            return i*n + j

        def valid(x,y):
            nonlocal n
            return 0<=x<n and 0<=y<n

        def unionNeighbors (x,y):
            nonlocal grid
            xyID = getID(x,y)
            for inc in [(-1,0), (0, -1), (1,0), (0,1)]:
                nX, nY = x+inc[0], y+inc[1]
                if valid(nX,nY) and grid[nX][nY]==1:
                    union(xyID, getID(nX,nY))

        def getUnionArea (x, y):
            nonlocal grid
            xyID = getID(x,y)
            ids  = set()
            area = 1
            for inc in [(-1,0), (0, -1), (1,0), (0,1)]:
                nX, nY = x+inc[0], y+inc[1]
                if valid(nX,nY) and grid[nX][nY]==1:
                    parentID = find(getID(nX,nY))
                    if parentID not in ids:
                        ids.add(parentID)
                        area += parent[parentID][1]
            #print(f"({x},{y}) Combine area of {ids} area={area}")
            return area

        n = len(grid)
        if n==1:
            return 1

        parent = [] #(parentID=x*n+y,area)
        idx    = 0
        # init parent to [id,area]
        for i in range(n):
            for j in range(n):
                parent.append([idx,grid[i][j]])
                idx += 1
        # go over do the union-find
        for i in range(n):
            for j in range(n): # union with above and left not working!!!
                if grid[i][j]==1:
                    unionNeighbors(i,j)
        #print(parent)
        maxArea = 0
        for i in range(n):
            for j in range(n): 
                if grid[i][j]==0: # switch to 1
                    maxArea=max(maxArea, getUnionArea(i,j))        
        return maxArea if maxArea>0 else n*n

        # parent = {} # id: parent_id
        # count  = defaultdict(int) # id: member_count
        # tops   = set() # group ids, it is in fact not used
        # for i in range(n):
        #     for j in range(n):
        #         if grid[i][j]==0:
        #             continue
        #         # grid[i][j]=1
        #         currId = getID(i,j)
        #         tops.add(currId)
        #         count[currId] = 1
        #         parent[currId] = currId
        #         if j>0 and grid[i][j-1]==1: 
        #             union(currId, getID(i,j-1)) # union with left
        #         if i>0 and grid[i-1][j]==1:
        #             union(currId, getID(i-1,j)) # union with up

        # if len(tops)==0:
        #     return 1
        # largest = max([count[t] for t in tops])
        # #print(f"parent={parent}")
        # #print(f"tops={tops}")
        # #print(f"count={count}")
        # #print(f"Before largest={largest}")
        # #no need to union, just add up neighbor's groups
        # for i in range(n):
        #     for j in range(n):
        #         if grid[i][j]==1:
        #             continue
        #         # grid[i][j]=0, switch to 1 
        #         groups = set()
        #         for inc in [(-1,0),(1,0),(0,-1),(0,1)]:
        #             nI, nJ = i+inc[0], j+inc[1]
        #             if valid(nI,nJ) and grid[nI][nJ] == 1:
        #                 groups.add(find(getID(nI,nJ)))
        #         total = 1+sum([count[g] for g in groups])
        #         largest = max(largest, total)
        # return largest
58168