https://www.codewars.com/kata/52bb6539a4cf1b12d90005b7/train/python

battlefield_ship_validator.py

# island finding with disjoint set union https://judge.yosupo.jp/submission/113653
from collections import defaultdict, Counter
class DSU:
    def __init__(self, n):
        self.par = list(range(n))
        self.size = [1]*n
    def find(self, x):
        bkx = x
        while x != self.par[x]:
            x = self.par[x]
        while self.par[bkx] != x:
            self.par[bkx], bkx = x, self.par[bkx]
        return x
    def union(self, x, y):
        x,y = self.find(x), self.find(y)
        if x == y: return False
        if self.size[x] < self.size[y]: x,y=y,x
        self.par[y] = x
        self.size[x] += self.size[y]
        self.size[y] = 0
        return True
# Argument is guaranteed to be 10*10 two-dimension array
n = 10
DIRS = [(0,1),(1,0),(0,-1),(-1,0), # neither by edge
       (1,1),(-1,-1),(1,-1),(-1,1)] # nor by corner.
def validate_battlefield(field):
    def in_bounds(r,c):
        return 0<=r<n and 0<=c<n
    # write your magic here
    dsu = DSU(n*n)
    occupied_indices = []
    for r in range(n):
        for c in range(n):
            if field[r][c] == 0: continue
            uid = r * n + c
            occupied_indices.append(uid)
            for dr,dc in DIRS:
                nr,nc=r+dr,c+dc
                if not in_bounds(nr, nc) or field[nr][nc] == 0: continue
                nid = nr*n+nc
                dsu.union(uid, nid)
    d = defaultdict(list)
    for uid in occupied_indices:
        d[dsu.find(uid)].append(uid)
#     Each ship must be a straight line, except for submarines, which are just single cell
# There must be single battleship (size of 4 cells), 2 cruisers (size 3), 3 destroyers (size 2) and 4 submarines (size 1). Any additional ships are not allowed, as well as missing ships.
    types = Counter()
    for island in d.values():
        diffs = []
        for i in range(1,len(island)):
            diffs.append(island[i] - island[i-1])
        if not ( all(map(lambda x: x == 10, diffs)) or all(map(lambda x: x == 1, diffs)) ):
            return False
        types[len(island)]+=1
    if types[4] != 1 or types[3] != 2 or types[2] != 3 or types[1] != 4: return False
    return set(types.keys()) == {1,2,3,4} and True