Транспортная задача (ТПР)

import copy

import numpy as np

B = np.array([100, 140, 100, 60])
S = np.array([
    [5, 4, 3, 2],
    [2, 3, 5, 6],
    [3, 2, 4, 3],
    [4, 1, 2, 4]
])
A = np.array([120,
              60,
              80,
              140])


# B = np.array([25,10,20,30,15])
# S = np.array([
#     [5, 3, 4, 6, 4],
#     [3, 4, 10, 5, 7],
#     [4, 6, 9, 3, 4],
# ])
# A = np.array([40,
#               20,
#               40])
#

def prt(res):  # вывод массива
    global S
    k = 0
    for i in range(res.shape[0]):
        for j in range(res.shape[1]):
            if res[i][j] > 0:
                k += S[i][j] * res[i][j]  # стоимость * кол-во

    print(res, k)
    print("--------")
    print()
    return res, k


def sv(b=True):  # северо-западный метод
    global A, B, S
    res = np.zeros_like(S)
    for i in range(A.shape[0]):
        for j in range(B.shape[0]):
            x = min(B[j], A[i])  # минимум между тем, сколько необходимо и тем, сколько имеется
            B[j] -= x  # вычитание взятого
            A[i] -= x  # вычитание взятого
            res[i][j] += x  # сохранение результата
    if b:
        res[res == 0] = -1  # поставить на месте всех не использованных -1
    return prt(res)


def min_st(A, B, S):  # метод минимальной стоимости
    res = np.zeros_like(S)
    l = []
    for i in range(S.shape[0]):
        for j in range(S.shape[1]):
            l.append((S[i][j], i, j))  # сохранение цены и позиций

    l.sort()  # сортировка

    for _, i, j in l:
        x = min(B[j], A[i])  # минимум между тем, сколько необходимо и тем, сколько имеется
        B[j] -= x  # вычитание взятого
        A[i] -= x  # вычитание взятого
        res[i][j] += x  # сохранение результата
    return prt(res)


recs = []  # позиции элементов, образующих многоугольник


def rec(i, j, f, D, X, x, y, ot=[]):  # метод составление многоугольника
    global recs
    if len(recs) != 0:  # выход, если многоугольник составлен
        return
    if i >= X.shape[0] or i < 0 or j >= X.shape[1] or j < 0 or D[i][j]:  # выход, если
        # многоугольник составлен или выход за границы или уже посетили данную клетку
        return
    if i == x and j == y and len(ot) > 2:  # найден ответ
        recs = ot.copy()
        return
    D[i][j] = True  # отметка клетки
    if f:  # идем по вертикали
        rec(i + 1, j, f, D, X, x, y, ot)
        rec(i - 1, j, f, D, X, x, y, ot)
    else:  # идем по горизонтали
        rec(i, j + 1, f, D, X, x, y, ot)
        rec(i, j - 1, f, D, X, x, y, ot)
    b = False
    if X[i][j] != -1:
        ot.append((i, j))
        b = True
    if b:  # если на клетке стоит число
        if f:  # меняем направление
            rec(i, j + 1, not f, D, X, x, y, ot)
            rec(i, j - 1, not f, D, X, x, y, ot)
        else:  # меняем направление
            rec(i + 1, j, not f, D, X, x, y, ot)
            rec(i - 1, j, not f, D, X, x, y, ot)
        ot.pop()

    D[i][j] = 0  # ставим обратно, как непосещенную


def pt(X, res):  # метод перепросчета решения
    global S, recs
    print(X, int(res))
    print("********")
    recs = []
    u, v = np.zeros(X.shape[0]), np.zeros(X.shape[1])
    u[0] = 0
    l = []
    for i in range(X.shape[0]):
        for j in range(X.shape[1]):
            if X[i][j] >= 0:  # сохранение ненулевых элементов
                l.append((i, j))
    k1, k2 = {0}, set()  # u, v которые уже найдены
    while len(k1) + len(k2) < X.shape[0] + X.shape[1]:
        for a, b in l:
            if a in k1:
                v[b] = S[a][b] - u[a]  # вычисление нового вертикального элемента
                k2.add(b)
            if b in k2:
                u[a] = S[a][b] - v[b]  # вычисление нового диагонального элемента
                k1.add(a)
    m = []
    for i in range(X.shape[0]):
        for j in range(X.shape[1]):
            if X[i][j] < 0:
                m.append((S[i][j] - (u[i] + v[j]), i, j))  # вычисление относительных оценок
    m.sort()  # сортировка относительных оценок
    if m[0][0] > 0:  # ответ найден
        return
    rec(m[0][1] + 1, m[0][2], True, np.zeros_like(X), X, m[0][1], m[0][2])  # составляем многоугольник (по диагонали)
    rec(m[0][1] - 1, m[0][2], True, np.zeros_like(X), X, m[0][1], m[0][2])
    rec(m[0][1], m[0][2] + 1, False, np.zeros_like(X), X, m[0][1], m[0][2])  # составляем многоугольник (по вертикали)
    rec(m[0][1], m[0][2] - 1, False, np.zeros_like(X), X, m[0][1], m[0][2])
    mi = np.min([X[recs[i][0]][recs[i][1]] for i in range(len(recs)) if i % 2 == 0])  # минимум в многоугольнике
    b = True
    for i in range(len(recs)):
        if i % 2 == 0:
            X[recs[i][0]][recs[i][1]] -= mi  # позиция на минусе => вычитание
            if X[recs[i][0]][recs[i][1]] == 0 and b:
                b = False
                X[recs[i][0]][recs[i][1]] = -1  # удаление первого 0
        else:
            X[recs[i][0]][recs[i][1]] += mi  # позиция на плюсе => прибавление
    X[m[0][1]][m[0][2]] = mi  # добавление новой вершины
    res += m[0][0] * mi  # пересчитывание результата
    pt(X, res)  # запуск с новой таблицей

min_st(copy.deepcopy(A), copy.deepcopy(B), copy.deepcopy(S))
pt(*sv())