Skip to content

Instantly share code, notes, and snippets.

Embed
What would you like to do?
import itertools
import networkx as nx
""" DATA """
S = ('S', [0, 1, 2]) # ranking: good to bad
K = ('K', [1, 2, 0])
L = ('L', [1, 2, 0])
D = ('D', [0, 2, 1])
A = ('A', [0, 2, 1])
J = ('J', [2, 1, 0])
P = [S, K, L, D, A, J]
C = [('c0', 0, 1), ('c1', 1, 2), ('c2', 2, 3)] # (name, index, limit)
""" Helper """
def get_cost(ranking, index):
return ranking.index(index)**2 # quadratic penalty
""" Min-cost flow graph """
G = nx.DiGraph()
# Demands
for p in P:
G.add_node(p[0], demand=-1)
G.add_node('sink', demand=len(P))
# layer 1
for p, c in itertools.product(P, C):
G.add_edge(p[0], c[0], weight=get_cost(p[1], c[1]))
# layer 2
for c in C:
G.add_edge(c[0], 'sink', capacity=c[2])
flowCost, flowDict = nx.network_simplex(G)
print('Cost: ', flowCost)
""" SOLUTION """
for p in P:
for c in C:
if flowDict[p[0]][c[0]]:
print(p[0], '->', c[0])
# Output:
# -------
# Cost: 2
# S -> c0
# K -> c1
# L -> c1
# D -> c2
# A -> c2
# J -> c2
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment