ВИ Лаб 2: Формирање на најдобар тим (CSP)

from constraint import *


def get_key(val, the_dict):
    for k, v in the_dict.items():
        if val == v:
            return k

    return "no such key"


if __name__ == '__main__':

    num_members = int(input())
    members_dict = dict()
    members_domain = list()
    members_vars = [f'Participant {member + 1}' for member in range(5)]
    for _ in range(num_members):
        w, n = input().split(" ")
        members_dict[n] = float(w)
        members_domain.append(float(w))

    num_leaders = int(input())
    leaders_dict = dict()
    leaders_domain = list()
    leaders_vars = ["Leader"]
    for _ in range(num_leaders):
        w, n = input().split(" ")
        leaders_dict[n] = float(w)
        leaders_domain.append(float(w))

    problem = Problem()

    problem.addVariables(leaders_vars, leaders_domain)

    problem.addVariables(members_vars, members_domain)

    problem.addConstraint(AllDifferentConstraint(), members_vars)

    problem.addConstraint(MaxSumConstraint(100))

    solutions = dict()

    for solution in problem.getSolutions():
        total_weight = 0
        for value in solution.values():
            total_weight += value

        solutions[total_weight] = [value for value in solution.values()]

    best_team = max(solutions.keys())

    whole_team = solutions[best_team]

    best_leader = whole_team[-1]

    print(f'Total score: {round(best_team, 1)}')
    print(f'Team leader: {get_key(best_leader, leaders_dict)}')

    del whole_team[-1]

    for i in range(len(whole_team)):
        print(f'Participant {i + 1}: {get_key(whole_team[i], members_dict)}')

    ''' ---------- Sample input and ouput ----------
    Input       Result
    10          Total score: 100.0
    31.3 A      Team leader: O
    28.4 B      Participant 1: J
    26.1 C      Participant 2: I
    24.2 D      Participant 2: I
    21.8 E      Participant 3: H
    20.3 F      Participant 4: F
    15.5 G      Participant 5: B
    14.1 H
    12.5 I
    11.5 J
    5
    32.2 K
    27.4 L
    24.6 M
    14.9 N
    13.2 O
    '''