team_division.py

team_division.py

"""
# team_division.py
n人をkチーム（各チームはn//kまたはn//k+1人）に分けることを複数回実施する場面で、なるべく効率的に、参加者が自分以外の全員と同じチームになれるようなチームの分け方を出力するプログラムです。

# 使い方

```
pip install tqdm
```

```
team_division.py [n] [k]
```

- n: メンバーの総数
- k: チーム数

出力は１つめがチームの分け方、２つめが最終的に各人が同じチームになった回数のマトリクス。
"""
import tqdm
from itertools import combinations

def generate_team_combinations(n, k):
    # 各チームの最小人数
    min_team_size = n // k
    # 余りの人数
    remainder = n % k

    # チームのサイズを決定
    team_sizes = [min_team_size + 1 if i < remainder else min_team_size for i in range(k)]

    # 全ての組み合わせを生成
    def generate_combinations(people, sizes):
        if not sizes:
            return [[]]
        current_size = sizes[0]
        rest_sizes = sizes[1:]
        combinations_list = []
        for team in combinations(people, current_size):
            remaining_people = [p for p in people if p not in team]
            for rest in generate_combinations(remaining_people, rest_sizes):
                combinations_list.append([list(team)] + rest)
        return combinations_list

    # 人のリストを作成
    people = list(range(n))
    return generate_combinations(people, team_sizes)

def find_min_cost_combination(team_combinations, history_matrix):
    min_cost = float('inf')
    best_combination = None

    for combination in team_combinations:
        cost = 0
        for team in combination:
            for i in range(len(team)):
                for j in range(i + 1, len(team)):
                    cost += history_matrix[team[i]][team[j]]
        
        if cost < min_cost:
            min_cost = cost
            best_combination = combination

    return best_combination
  
def update_history_matrix(history_matrix, teams):
    updated_history_matrix = [row[:] for row in history_matrix]
    for team in teams:
        for i in range(len(team)):
            for j in range(i + 1, len(team)):
                updated_history_matrix[team[i]][team[j]] += 1
                updated_history_matrix[team[j]][team[i]] += 1
    return updated_history_matrix
  
def generate_teams(n, k, max_iter=20):
    # team_combinationsをすべて求める
    team_combinations = generate_team_combinations(n, k)
    
    # history_matrixの初期値（すべて0）を生成
    history_matrix = [[0] * n for _ in range(n)]
    # 対角成分は1にする
    for i in range(n):
        history_matrix[i][i] = 1
    
    # 結果のリスト
    result = []
    
    for _ in tqdm.tqdm(range(max_iter)): #最大20回繰り返す
        # find_min_cost_combinationでbest_teamを求め、結果のリストに追加
        square_history_matrix = [[value ** 2 for value in row] for row in history_matrix]
        best_team = find_min_cost_combination(team_combinations, square_history_matrix)
        result.append(best_team)
        
        # update_history_matrixでhistory_matrixを更新する
        history_matrix = update_history_matrix(history_matrix, best_team)
        
        # もしhistory_matrixのすべての要素が1以上ならbreak
        if all(all(value >= 1 for value in row) for row in history_matrix):
            break
    
    return result, history_matrix

if __name__ == "__main__":
    import argparse

    parser = argparse.ArgumentParser(description='Generate teams based on given parameters.')
    parser.add_argument('n', type=int, help='The total number of players')
    parser.add_argument('k', type=int, help='The number of teams')
    parser.add_argument('--max_iter', type=int, default=20, help='The maximum number of iterations')
    args = parser.parse_args()

    result, history_matrix = generate_teams(args.n, args.k, args.max_iter)
    print("result")
    for i, team in enumerate(result):
        print(f"{i}: {team}")
    print("history_matrix")
    print(history_matrix)