20204883 multicast routing.py

20204883 multicast routing.py

from ortools.sat.python import cp_model
from copy import deepcopy


class VarArraySolutionPrinter(cp_model.CpSolverSolutionCallback):
    """Print intermediate solutions."""

    def __init__(self, variables):
        cp_model.CpSolverSolutionCallback.__init__(self)
        self.__variables = variables
        self.__solution_count = 0

    def on_solution_callback(self):
        self.__solution_count += 1
        for v in self.__variables:
            print('%s=%i' % (v, self.Value(v)), end=' ')
        
        print(f'\nObjectiveValue={self.ObjectiveValue()}')

    def solution_count(self):
        return self.__solution_count

'''
def print_matrix(a):
    for row in a:
        for element in row:
            print(f'{str(element):8s}', end=' ')
        print()
    print()
'''

# read inp
file_name = 'multicast.txt'
with open(file_name, 'r') as f:
    #              start_node, max_time
    node_count, edge_count, s, L = map(int, f.readline().split())
    s -= 1

    V = list(range(node_count))

    # cost, time
    c = [[None] * node_count for _ in V]
    t = [[None] * node_count for _ in V]
    for _ in range(edge_count):
        i, j, t_edge, c_edge = map(int, f.readline().split())
        i -= 1
        j -= 1
        t[i][j] = t[j][i] = t_edge
        c[i][j] = c[j][i] = c_edge

# print_matrix(t)
# print_matrix(c)


# model
model = cp_model.CpModel()


# vars
x = [[model.NewIntVar(0, 1, f'x[{i}][{j}]')
        for j in V]
        for i in V]

max_total_time = sum([sum(filter(lambda x: bool(x), row)) for row in t])
y = [model.NewIntVar(0, max_total_time, f'y[{i}]')
        for i in V]


# sets 
V_squared = [(i, j) for i in V for j in V]
E = list(filter(lambda tup: t[tup[0]][tup[1]], V_squared))
# print(E)
# print(len(E))

V_no_s = deepcopy(V)
V_no_s.remove(s)

A = [[j for j in V if (i,j) in E] for i in V]
# print(A)


# constraints
for i in V_no_s:
    model.Add(sum([x[j][i] for j in A[i]]) == 1)

model.Add(sum([x[s][i] for i in A[s]]) >= 1)

for i, j in E:
    b = model.NewBoolVar('b')
    model.Add(x[i][j] == 1).OnlyEnforceIf(b)
    model.Add(x[i][j] == 0).OnlyEnforceIf(b.Not())
    model.Add(y[j] == y[i] + t[i][j]).OnlyEnforceIf(b)

    model.Add(y[i] <= L)

'''
x[0][0]=0 x[0][1]=1 x[0][2]=1 x[0][3]=1 x[0][4]=0 x[0][5]=0 x[0][6]=0 x[1][0]=0 x[1][1]=0 x[1][2]=0 x[1][3]=0 x[1][4]=0 x[1][5]=0 x[1][6]=1 x[2][0]=0 x[2][1]=0 x[2][2]=0 x[2][3]=0 x[2][4]=0 x[2][5]=0 x[2][6]=0 x[3][0]=0 x[3][1]=0 x[3][2]=0 x[3][3]=0 x[3][4]=1 x[3][5]=0 x[3][6]=0 x[4][0]=0 x[4][1]=0 x[4][2]=0 x[4][3]=0 x[4][4]=0 x[4][5]=1 x[4][6]=0 x[5][0]=0 x[5][1]=0 x[5][2]=0 x[5][3]=0 x[5][4]=0 x[5][5]=0 x[5][6]=0 x[6][0]=0 x[6][1]=0 x[6][2]=0 x[6][3]=0 x[6][4]=0 x[6][5]=0 x[6][6]=0 29.0
Add a constraint: no route comes to s
'''
for i in A[s]:
    model.Add(x[i][s] == 0)


# objective
for i, j in E:
    objective_function = sum([c[i][j] * x[i][j] for i, j in E])
model.Minimize(objective_function)


# solve
solution_printer = VarArraySolutionPrinter([x[i][j] for i in V for j in V])
cp_model.CpSolver().Solve(model, solution_callback=solution_printer)