Untitled

def solution(matrix):
    def is_valid_diagonal(x, y, dx, dy):
        cx, cy, i = x, y, 0

        # track first num is 2 or 0
        pattern = [2, 0]
        first_num = 0
        while 0 <= cx + dx < len(matrix) and 0 <= cy + dy < len(matrix[0]):
            cx += dx
            cy += dy

            if i == 0:
                first_num = pattern.index(matrix[cx][cy])

            if i % 2 == 0 and matrix[cx][cy] != pattern[first_num]:
                return False
            elif i % 2 == 1 and matrix[cx][cy] != pattern[first_num + 1 % 2]:
                return False

            i += 1

        # diagonal is valid if 2, 0 is in line
        return i > 0


    q = []
    dirs = [(1, 1), (-1, -1), (-1, 1), (1, -1)]
    diagonals = 0

    # get all initial positions of 1
    for row in range(len(matrix)):
        for col in range(len(matrix[0])):
            if matrix[row][col] == 1:
                q.append((row, col))

    # for each 1 in the matrix
    while q:
        x, y = q.pop()

        # get all diagonals at 1 starting positions
        for dx, dy in dirs:
            # count valid diagonals
            if is_valid_diagonal(x, y, dx, dy):
                diagonals += 1

    return diagonals

print(solution(matrix))