Transportation problem solver in Python

transport.py

import numpy as np
from collections import Counter


def transport(supply, demand, costs):
    assert sum(supply) == sum(demand)

    s = np.copy(supply)
    d = np.copy(demand)
    C = np.copy(costs)

    n, m = C.shape

    X = np.zeros((n, m))
    indices = [(i, j) for i in range(n) for j in range(m)]
    xs = sorted(zip(indices, C.flatten()), key=lambda a_b: a_b[1])

    for (i, j), _ in xs:
        if d[j] == 0:
            continue
        else:
            remains = s[i] - d[j] if s[i] >= d[j] else 0
            grabbed = s[i] - remains
            X[i, j] = grabbed
            s[i] = remains
            d[j] -= grabbed

    while True:
        u = np.array([np.nan] * n)
        v = np.array([np.nan] * m)
        S = np.zeros((n, m))

        _x, _y = np.where(X > 0)
        nonzero = list(zip(_x, _y))
        f = nonzero[0][0]
        u[f] = 0

        while any(np.isnan(u)) or any(np.isnan(v)):
            for i, j in nonzero:
                if np.isnan(u[i]) and not np.isnan(v[j]):
                    u[i] = C[i, j] - v[j]
                elif not np.isnan(u[i]) and np.isnan(v[j]):
                    v[j] = C[i, j] - u[i]
                else:
                    continue

        for i in range(n):
            for j in range(m):
                S[i, j] = C[i, j] - u[i] - v[j]

        s = np.min(S)
        if s >= 0:
            break

        i, j = np.argwhere(S == s)[0]
        start = (i, j)

        T = np.copy(X)
        T[start] = 1
        while True:
            _xs, _ys = np.nonzero(T)
            xcount, ycount = Counter(_xs), Counter(_ys)

            for x, count in list(xcount.items()):
                if count <= 1:
                    T[x, :] = 0
            for y, count in list(ycount.items()):
                if count <= 1:
                    T[:, y] = 0

            if all(x > 1 for x in list(xcount.values())) \
                    and all(y > 1 for y in list(ycount.values())):
                break

        dist = lambda p1, p2: abs(p1[0] - p1[1]) + abs(p2[0] - p2[1])
        fringe = set(tuple(p) for p in np.argwhere(T > 0))

        size = len(fringe)

        path = [start]
        while len(path) < size:
            last = path[-1]
            if last in fringe:
                fringe.remove(last)
            next = min(fringe, key=lambda x_y: dist(last, (x_y[0], x_y[1])))
            path.append(next)

        neg = path[1::2]
        pos = path[::2]
        q = min(X[tuple(zip(*neg))])
        X[tuple(zip(*neg))] -= q
        X[tuple(zip(*pos))] += q

    return X, np.sum(X * C)


# Метод транспортных потенциалов
if __name__ == '__main__':
    # Наличие
    demand = np.array([30, 90, 50])
    # Потребность
    supply = np.array([70, 60, 30, 10])
    # Тариф
    costs = np.array([
        [8, 3, 1],
        [4, 7, 4],
        [5, 2, 6],
        [0, 0, 0],
    ])

    print(transport(supply, demand, costs))