Challenge_2 | Python Fiddle

#This is essentially a DFS problem but want to solve it non recursively
def solution(mat):
    #generate all the points that have '1' in them, representing an individual stream
    pts = [(i,j) for i in range(0,len(mat)) for j in range(0,len(mat[1])) if mat[i][j]==1]
    streams = [] #contains list of points that make each stream
   
    #make sure all the points are visited, once each point is visited, remove it from the list of 
    while pts:
        temp = [pts[0]] #at each point
        del pts[0]
        
        neighbourFound = True
        while neighbourFound:
            neighbourFound = False #if at the end of the loop no neighbour is found we proceed to look for a new stream
            for i in range(len(pts)): #check for neighbours among all remaining points and those that are already part of a stream
                for j in range(len(temp)):
                    if i>=len(pts): # loop mutates pts so avoid i from exceeding actual current length of pts
                        break
                    if isneighbour(pts[i],temp[j],mat):
                        neighbourFound = True
                        temp.append(pts[i]) #add the point to a temporary stream
                        del pts[i] #remove the point from pts and add to a stream
        streams.append(temp)
    
    return [len(x) for x in streams]
    
    
#check if two points in a matrix can form a stream or in other words are neighbours according the question
def isneighbour(pt1,pt2,mat):
    #simple trick if they are diagonally, vertically or horizontal neigbours then they are within the 1-neighbourhood of each other
    if abs(pt1[0]-pt2[0])<=1 and abs(pt1[1]-pt2[1])<=1:
        return True
    #if they are points at the edges and wrap around horizontally
    elif pt1[0]==pt2[0]:
        if abs(pt1[1]-pt2[1])==len(mat[0])-1:
            return True
    #if they are points at the edges and wrap around vertically
    elif pt1[1] ==pt2[1]:
        if abs(pt1[0]-pt2[0]) == len(mat)-1:
            return True
    #if they are points at the edges and wrap around diagonally    
    elif abs(pt1[0]-pt2[0])== len(mat) - 1 and abs(pt1[1]-pt2[1])== len(mat[0])-1:
        return True
    return False
 
mat = [
[1, 0, 0, 0, 0],
[0, 0, 1, 0, 0],
[0, 0, 0, 0, 1]
]
print(solution(mat))

Python Fiddle

Python Cloud IDE