Vogel algorithm for transportation problem

vogel_approximation.py

import math

supply = [50, 60, 50, 50]
demand = [30, 20, 70, 30, 60]

costs = [
    [16, 16, 13, 22, 17],
    [14, 14, 13, 19, 15],
    [19, 19, 20, 23, 50],
    [50, 12, 50, 15, 11]
]

nRows = len(supply)
nCols = len(demand)

rowDone = [False] * nRows
colDone = [False] * nCols
results = [[0 for _ in range(nCols)] for _ in range(nRows)]

def next_cell():
    res1 = max_penalty(nRows, nCols, True)
    res2 = max_penalty(nCols, nRows, False)
    if res1[3] == res2[3]:
        return res1 if res1[2] < res2[2] else res2
    return res2 if res1[3] > res2[3] else res1

def diff(j, l, is_row):
    min1 = math.inf
    min2 = math.inf
    minP = -1
    for i in range(l):
        done = colDone[i] if is_row else rowDone[i]
        if done:
            continue
        c = costs[j][i] if is_row else costs[i][j]
        if c < min1:
            min2, min1, minP = min1, c, i
        elif c < min2:
            min2 = c
    return [min2 - min1, min1, minP]

def max_penalty(len1, len2, is_row):
    md = -math.inf
    pc, pm, mc = -1, -1, -1
    for i in range(len1):
        done = rowDone[i] if is_row else colDone[i]
        if done:
            continue
        res = diff(i, len2, is_row)
        if res[0] > md:
            md, pm, mc, pc = res[0], i, res[1], res[2]
    return [pm, pc, mc, md] if is_row else [pc, pm, mc, md]

def main():
    supply_left = sum(supply)
    total_cost = 0
    while supply_left > 0:
        cell = next_cell()
        r, c = cell[0], cell[1]
        q = min(demand[c], supply[r])
        demand[c] -= q
        colDone[c] = demand[c] == 0
        supply[r] -= q
        rowDone[r] = supply[r] == 0
        results[r][c] = q
        supply_left -= q
        total_cost += q * costs[r][c]

    print("    A   B   C   D   E")
    for i, result in enumerate(results):
        print(chr(ord('W') + i), ' '.join(f"{item:3d}" for item in result))
    print("\nTotal cost =", total_cost)

if __name__ == "__main__":
    main()