dfs.py

"""
Дан неориентированный невзвешенный граф.
Для него вам необходимо найти количество вершин, лежащих в одной компоненте связности с данной вершиной (считая эту вершину).

В первой строке входных данных содержатся два числа: N и S (1 ≤ N ≤ 100; 1 ≤ S ≤ N), где N – количество вершин графа, а S – заданная вершина.
В следующих N строках записано по N чисел – матрица смежности графа, в которой 0 означает отсутствие ребра между вершинами, а 1 – его наличие.
Гарантируется, что на главной диагонали матрицы всегда стоят нули.

Выведите одно целое число – искомое количество вершин.
"""


def solve():
    # Прочесть: N - количество вершин в графе, S - заданная вершина
    # пример: "4 1".split() -> ["4", "1"] -> list(map(int, ["4", "1"])) -> [4, 1]
    # n = 4, s = 1
    n, s = list(map(int, input().split()))

    # Прочесть матрицу смежности
    s -= 1
    A = []
    # Прочитать n строк
    for i in range(n):
        # Прочесть строку матрицы смежности
        # пример: "0 1 0 1".split() -> ["0", "1", "0", "1"] -> list(map(int, ["0", "1", "0", "1"])) -> [0, 1, 0, 1]
        A.append(list(map(int, input().split())))  # A[i][j] - ребро из вершины i в вершину j

    # mark - массив размера n, в котором будем хранить цвета вершин
    # пример: [0] * 3 -> [0, 0, 0], т.е. все вершины не раскрашены
    mark = [0] * n

    def dfs(k, color):
        """
        Поиск в глубину

        Рекурсивная функция, которая помечает все вершины, смежные с вершиной k, цветом color

        Пример:
        ------

            Граф:
            -----
            0 - 1  2
            |   
            3

            4 вершины и 2 ребра

            Матрица смежности:
            ------------------
            0 1 0 1  # 0 cмежна с 1 и 3
            1 0 0 0  # 1 cмежна с 0
            0 0 0 0  # 2 не cмежна ни с кем
            1 0 0 0  # 3 cмежна с 0

            n = 4, s = 0
                                         mark = [0, 0, 0, 0]
            dfs(0, 1), k = 0, color = 1, mark = [1, 0, 0, 0]
            dfs(1, 1), k = 1, color = 1, mark = [1, 1, 0, 0]
            dfs(3, 1), k = 3, color = 1, mark = [1, 1, 0, 1]

            Ответ: sum(mark) = 1 + 1 + 0 + 1 = 3

        k - номер вершины, с которой начинаем поиск
        color - цвет, которым будем помечать вершины, которые нашли
        """
        # Пометить вершину k цветом color
        mark[k] = color

        # Перебрать все вершины, смежные с k
        for i in range(n):
            # Если вершина i смежна с k и еще не помечена
            if A[k][i] == 1 and mark[i] == 0:
                # Пометить вершину i цветом color
                dfs(i, color)

    # Вызвать поиск в глубину из вершины s c цветом 1
    # Красим все смежные вершины цветом 1
    dfs(s, 1)

    # Вывести количество вершин, помеченных цветом 1 – это и будет ответом, т.е. количество вершин в компоненте связности c вершиной s
    print(sum(mark))