Создание цепи Маркова

Написал код, а распределение вероятностей через три шага равно 1 0 0 0 0 0. Не могу это исправить. Задание:

1. Задать цепь Маркова: множество состояний, начальное распределение (все случаи), матрицу переходных вероятностей, граф цепи Маркова.
2. Найти распределение состояний цепи Маркова на любом шаге: вход задаем число состояний, начальное распределение, число шагов.
3. Классифицировать состояния цепи Маркова. Вариант: Урна содержит N белых шаров. За 1 шаг из урны по схеме случайного равновероятного выбора вынимают один шар и заменяют его черным. Множество состояний – число белых шаров Мой код:

import numpy as np
import networkx as nx
import matplotlib.pyplot as plt

def create_transition_matrix(N):
    """
    Создает матрицу переходных вероятностей для задачи.
    """
    P = np.zeros((N+1, N+1))  # Матрица переходов размером (N+1) x (N+1)

    for k in range(N+1):
        if k > 0:  # Переход в состояние k-1 (вынут белый шар)
            P[k, k-1] = k / (k + (N-k))  # k белых из общего количества
        if k < N:  # Переход в то же состояние (вынут черный шар)
            P[k, k] = (N-k) / (k + (N-k))  # (N-k) черных из общего количества

    return P


def compute_distribution(P, pi0, steps):
    """
    Вычисляет распределение вероятностей на любом шаге.
    P: матрица переходных вероятностей.
    pi0: начальное распределение (вектор длины N+1).
    steps: число шагов.
    """
    pi_t = np.array(pi0)
    for _ in range(steps):
        pi_t = pi_t @ P  # Итеративное умножение для расчета распределения
    return pi_t

def classify_states(P):
    """
    Классифицирует состояния цепи Маркова.
    P: матрица переходных вероятностей.
    """
    N = P.shape[0]
    classifications = []
    for state in range(N):
        if np.all(P[:, state] == 0):  # Проверяем, поглощающее ли состояние
            classifications.append((state, "Поглощающее"))
        else:
            classifications.append((state, "Переходное"))
    return classifications

def draw_graph(P):
    """
    Строит граф цепи Маркова.
    P: матрица переходных вероятностей.
    """
    G = nx.DiGraph()
    N = P.shape[0]

    for i in range(N):
        for j in range(N):
            if P[i, j] > 0:
                G.add_edge(i, j, weight=round(P[i, j], 2))

    pos = nx.spring_layout(G)
    nx.draw(G, pos, with_labels=True, node_size=700, node_color="lightblue", font_size=10, font_weight="bold")
    edge_labels = nx.get_edge_attributes(G, 'weight')
    nx.draw_networkx_edge_labels(G, pos, edge_labels=edge_labels)
    plt.title("Граф цепи Маркова")
    plt.show()

# Параметры задачи
N = 5  # Число белых шаров в начале
P = create_transition_matrix(N)  # Матрица переходов

# 1. Матрица переходов и граф цепи Маркова
print("Матрица переходов P:")
print(P)

# Построим граф
draw_graph(P)

# 2. Распределение состояний
pi0 = [1] + [0] * N  # Начальное распределение: 100% белых шаров
steps = 3  # Число шагов
pi_t = compute_distribution(P, pi0, steps)

print(f"\nРаспределение вероятностей через {steps} шага(ов):")
print(pi_t)

# 3. Классификация состояний
classifications = classify_states(P)
print("\nКлассификация состояний:")
for state, class_type in classifications:
    print(f"Состояние {state}: {class_type}")