【AtCoder】ABC310 – D問題について解説【Python】

abc310-d プログラミング
2023.07.17

そまちょブログのそまちょ（@somachob）です。

この記事は、AtCoder Beginner Contest 309 の D問題についての解説です。

D – Peaceful Teams

問題へのリンク

入力例１をもとに解説します。

入力例１
5 2 2
1 3
3 4

入力を受け取ります。

A, B については NG というリストで管理しています。

選手の番号は -1 していることに注意してください。

# 入力
N, T, M = map(int, input().split())
NG = [[] for _ in range(N)] 
for _ in range(M):
    A, B = map(int, input().split())
    A -= 1
    B -= 1
    NG[A].append(B)
    NG[B].append(A)

# 入力確認
print(N, T, M)
print(NG)

""" 出力結果
5 2 2
[[2], [], [0, 3], [2], []]
"""

たとえば、NG[0] の値は、選手１の相性の悪い選手を表します。

NG[0] = 2 なので、選手１は選手３（2+1）と相性が悪いということです。

NG[2] = [0, 3] なので、選手３は【選手１と選手４】と相性が悪いことを表します。

入力の受け取りについてわからないときは、入力について解説している記事があるのでぜひ参考にしてください。

【Python】input の使い方

チーム分けを列挙する

この問題で、一番のポイントがチーム分けの実装です。

再帰関数でチーム分けを実現してみます。

再帰関数の基本は、次のようなコードになります。

def 関数名(引数):
    # ベースケース
    if 条件:
        return 値

    # 再帰ステップ
    else:
        return 関数名(引数を更新した値)

再帰関数について解説した記事も参考にしてください。

超初心者でもわかる再帰関数の基礎！競プロ初心者のためのPython実装ガイド

これをもとに各選手をチーム分けする関数を作成します。

各選手については、すでにあるいずれかのチームに所属させるか、新しいチームに所属させるかのパターンで実装します。

ベースケースは「すべての選手について処理したとき」として設定します。

これをするコードが次のコードです。N=3 のときの出力結果を表示させています。

def rec(player, teams):
    # ベースケース
    if player == N:
        # できたチームの確認
        print(teams)
        return

    # 再帰ステップ
    # すでにあるチームに所属させる
    for i in range(len(teams)):
        teams[i].append(player)	# チームに所属させる
        rec(player + 1, teams)	# 次の選手について再帰処理
        teams[i].pop()			# チームから抜ける

    # 新しいチームに所属させる
    teams.append([player])	# チームに所属させる
    rec(player + 1, teams)	# 次の選手について再帰処理
    teams.pop()				# チームから抜ける


rec(0,[])


"""出力結果
[[0, 1, 2]]
[[0, 1], [2]]
[[0, 2], [1]]
[[0], [1, 2]]
[[0], [1], [2]]
"""

全員が１つのチーム所属する場合からそれぞれが別チームに所属する場合まで出力できているのが確認できます。

このコードでチーム分けができるようになったので、「チーム数の判定」と「相性の悪さの判定」です。

できたチームの数を判定する

できたチーム数が T チームになっているかは、ベースケースで判定すればいいです。

N=５、T=2 のときの出力結果を表示させています。

def rec(player, teams):
    # ベースケース
    if player == N:
        # チーム数が T でないとき
        if len(teams) != T:
            return 
        
        # できたチームの確認
        print(teams)
        return

      
    # 再帰ステップ
    # すでにあるチームに所属させる
    for i in range(len(teams)):
        teams[i].append(player)	# チームに所属させる
        rec(player + 1, teams)	# 次の選手について再帰処理
        teams[i].pop()			# チームから抜ける

    # 新しいチームに所属させる
    teams.append([player])	# チームに所属させる
    rec(player + 1, teams)	# 次の選手について再帰処理
    teams.pop()				# チームから抜ける


rec(0,[])


"""出力結果
[[0, 1, 2, 3], [4]]
[[0, 1, 2, 4], [3]]
[[0, 1, 2], [3, 4]]
[[0, 1, 3, 4], [2]]
[[0, 1, 3], [2, 4]]
[[0, 1, 4], [2, 3]]
[[0, 1], [2, 3, 4]]
[[0, 2, 3, 4], [1]]
[[0, 2, 3], [1, 4]]
[[0, 2, 4], [1, 3]]
[[0, 2], [1, 3, 4]]
[[0, 3, 4], [1, 2]]
[[0, 3], [1, 2, 4]]
[[0, 4], [1, 2, 3]]
[[0], [1, 2, 3, 4]]
"""

２チームに分けたときのパターンだけ出力されているのが確認できます。

相性の悪い選手がいるか判定する

チームに相性の悪いペアがいないように定する必要があります。

選手を参加させようとしているチームに相性が悪い選手がいるかを判定します。

次の関数は、選手と所属させたいチームを受け取り、チーム内に相性が悪い選手がいるときは、True を返します。

# 相性の悪い選手がチームにいるか判定する関数
def is_NG(player, team):
    for member in team:
        # チームのメンバーと相性が悪いとき
        if member in NG[player]:
            return True

    return False

この関数を使って、選手をチームに所属させる前に相性の悪い選手がいるかを確認するようにします。

新しいチームに所属させるときは、1人だけなので、確認する必要はありません。

実装したコードです。入力例１での出力結果は、出力例１で示されているチームになっていることが確認できます。

# 入力
N, T, M = map(int, input().split())
NG = [[] for _ in range(N)] 
for _ in range(M):
    A, B = map(int, input().split())
    A -= 1
    B -= 1
    NG[A].append(B)
    NG[B].append(A)

    
# 相性の悪い選手がチームにいるか判定する関数
def is_NG(player, team):
    for member in team:
        # チームのメンバーと相性が悪いとき
        if member in NG[player]:
            return True

    return False
  

def rec(player, teams):
    # ベースケース
    if player == N:
        # チーム数が T でないとき
        if len(teams) != T:
            return 
        
        # できたチームの確認
        print(teams)
        return

      
    # 再帰ステップ
    # すでにあるチームに所属させる
    for i in range(len(teams)):
        # 相性の悪い選手がいるとき
        if is_NG(player, teams[i]):
            continue
      
        teams[i].append(player)	# チームに所属させる
        rec(player + 1, teams)	# 次の選手について再帰処理
        teams[i].pop()			# チームから抜ける

    # 新しいチームに所属させる
    teams.append([player])	# チームに所属させる
    rec(player + 1, teams)	# 次の選手について再帰処理
    teams.pop()				# チームから抜ける


rec(0,[])


"""出力結果
[[0, 1, 3, 4], [2]]
[[0, 1, 3], [2, 4]]
[[0, 3, 4], [1, 2]]
[[0, 3], [1, 2, 4]]
"""

チーム分けが何通りあるかを数える

条件を満たすチームではなく、条件を満たすチーム分けが何通りあるかを出力する必要があります。

答えを格納する変数 answer をグローバル宣言して、条件を満たすチームのときに １ 足しています。

次のコードで AC できます。

# 入力
N, T, M = map(int, input().split())
NG = [[] for _ in range(N)] 
for _ in range(M):
    A, B = map(int, input().split())
    A -= 1
    B -= 1
    NG[A].append(B)
    NG[B].append(A)


    
# 相性の悪い選手がチームにいるか判定する関数
def is_NG(player, team):
    for member in team:
        # チームのメンバーと相性が悪いとき
        if member in NG[player]:
            return True

    return False
  

def rec(player, teams):
    # グローバル宣言
    global answer
    
    # ベースケース
    if player == N:
        # チーム数が T でないとき
        if len(teams) != T:
            return 
        
        # 条件を満たすチームなので１足す
        answer += 1
        return

      
    # 再帰ステップ
    # すでにあるチームに所属させる
    for i in range(len(teams)):
        # 相性の悪い選手がいるとき
        if is_NG(player, teams[i]):
            continue
      
        teams[i].append(player)	# チームに所属させる
        rec(player + 1, teams)	# 次の選手について再帰処理
        teams[i].pop()			# チームから抜ける

    # 新しいチームに所属させる
    teams.append([player])	# チームに所属させる
    rec(player + 1, teams)	# 次の選手について再帰処理
    teams.pop()				# チームから抜ける

answer = 0
rec(0,[])
print(answer)

グローバル変数を使わずに実装すれば次のようになります。

イメージとすれば、すでにあるチームに所属させるときに条件を満たすチーム分けの数を数えて、その後、新しいチームに所属させたときに条件を満たすチーム分けの数を数えて、その数を返しているような感じです。

# 入力
N, T, M = map(int, input().split())
NG = [[] for _ in range(N)] 
for _ in range(M):
    A, B = map(int, input().split())
    A -= 1
    B -= 1
    NG[A].append(B)
    NG[B].append(A)


    
# 相性の悪い選手がチームにいるか判定する関数
def is_NG(player, team):
    for member in team:
        # チームのメンバーと相性が悪いとき
        if member in NG[player]:
            return True

    return False
  

def rec(player, teams):
    # 答えの変数
    answer = 0
    
    # ベースケース
    if player == N:
        # チーム数が T でないとき
        if len(teams) != T:
            return 0
        
        # 条件を満たすチーム
        return 1

      
    # 再帰ステップ
    # すでにあるチームに所属させる
    for i in range(len(teams)):
        # 相性の悪い選手がいるとき
        if is_NG(player, teams[i]):
            continue
      
        teams[i].append(player)				# チームに所属させる
        answer += rec(player + 1, teams)	# 次の選手について再帰処理
        teams[i].pop()						# チームから抜ける

    # 新しいチームに所属させる
    teams.append([player])				# チームに所属させる
    answer += rec(player + 1, teams)	# 次の選手について再帰処理
    teams.pop()							# チームから抜ける
    
    return answer

  
print(rec(0,[]))

参考

タイトルとURLをコピーしました