[백준] 14889. 스타트와 링크 (파이썬)

문제 : https://www.acmicpc.net/problem/14889

풀이

• 기본적으로 팀을 구성할 수 있는 모든 조합에 대해 두 팀의 능력치 차이를 구하는 브루트포스 알고리즘이다.

1. 변수설정

• N, S : input으로 받는 변수이다. N은 총 참가자의 수, S는 능력치 조합을 나타낸 2차원 배열이다.

• participant_list : 0 부터 N-1 까지의 값이 담긴 배열이다. 참가자 번호를 의미한다.

• answer : 답으로 출력된다. 최소값을 구해야하므로 INF, 2147000000 등 가능한 높은 수로 한다.

2. 조합

itertools 내장 모듈의 combinations 함수를 이용하면 조합을 손쉽게 출력할 수 있다.

이를 team_lists에 tuple 형태로 저장한다.

from itertools import combinations

participant_list = list(range(N))

team_lists = tuple(combinations(participant_list, N//2))
print(team_lists)

# N = 4일 경우 ((0, 1), (0, 2), (0, 3), (1, 2), (1, 3), (2, 3))
# N = 6일 경우 ((0, 1, 2), (0, 1, 3), (0, 1, 4), (0, 1, 5), (0, 2, 3), (0, 2, 4), (0, 2, 5), 
# (0, 3, 4), (0, 3, 5), (0, 4, 5), (1, 2, 3), (1, 2, 4), (1, 2, 5), (1, 3, 4), (1, 3, 5),
# (1, 4, 5), (2, 3, 4), (2, 3, 5), (2, 4, 5), (3, 4, 5))

tuple 자료형으로 형변환 하는 이유는 len 함수를 사용하기 위함이다.

3. 반복문

시간 단축의 핵심이 되는 부분이다.

team_lists = ((0, 1), (0, 2), (0, 3), (1, 2), (1, 3), (2, 3))
team_lists[0] 의 나머지는 team_lists[5]에 해당한다.
team_lists[1] 의 나머지는 team_lists[4]에 해당한다.
이런 식으로 team_lists[i]의 나머지는 team_lists[len(team_lists)-i-1]과 같다.

위 원리를 이용하여 한 팀은 team_lists[i], 나머지 팀은 team_lists[len(team_lists)-i-1]가 되도록 반복문을 돌면서 두 팀간 능력치 차이를 구하는 함수를 호출한다.

이렇게 구한 팀은 각각 team_start, team_link에 저장한다.

    for i in range(len(team_lists)//2):
        team_start = team_lists[i]
        team_link = team_lists[len(team_lists)-i-1]
        
        diff = get_diff(team_start, team_link)

4. 차이를 구하는 함수

i와 j가 한팀일 때 팀에 더해지는 능력치는 S[i][j] + S[j][i] 와 같다.
여기서 i와 j는 앞서 반복문에서 저장한 team_start, team_link에 담겨있다.

팀의 총 능력치는 팀에서 2명(i, j)을 뽑는 모든 조합으로 S[i][j] + S[j][i]를 모두 더하면 된다.
총 능력치를 저장할 변수를 0으로 초기화하고, 모든 조합을 반복문으로 돌면서 점수를 더하자.

똑같이 combinations 함수를 이용할 수 있지만, 2중 for문을 이용하는 것이 조금 더 편하다. (아래 코드 참조)

두 팀의 총 능력치를 구했다면, 그 차이의 절대값을 반환하면 된다.

def get_diff(team_start, team_link):
    stats_start = 0
    stats_link = 0
    length = len(team_start)

    for i in range(length):
        for j in range(i+1, length):
            stats_start += S[team_start[i]][team_start[j]] + S[team_start[j]][team_start[i]]
            stats_link += S[team_link[i]][team_link[j]] + S[team_link[j]][team_link[i]]
    return abs(stats_start - stats_link)

더하는 부분이 너무 길어서 지저분하다면 짧은 변수 이름이나, 2번으로 나눠서 더하는 게 좋을 것 같다.

출력

get_diff 함수를 통해 구한 값과 현재 answer의 값을 비교하여 최소값을 저장하면 된다.

diff = get_diff(team_start, team_link)
answer = min(answer, diff)

코드 (1152ms)

import sys
from itertools import combinations
input = sys.stdin.readline


def get_diff(team_start, team_link):
    stats_start = 0
    stats_link = 0
    length = len(team_start)

    for i in range(length):
        for j in range(i+1, length):
            stats_start += S[team_start[i]][team_start[j]] + S[team_start[j]][team_start[i]]
            stats_link += S[team_link[i]][team_link[j]] + S[team_link[j]][team_link[i]]
    return abs(stats_start - stats_link)


if __name__ == "__main__":
    N = int(input())
    S = [list(map(int, input().split())) for _ in range(N)]

    participant_list = list(range(N))
    answer = 2147000000
    team_lists = tuple(combinations(participant_list, N//2))

    for i in range(len(team_lists)//2):
        team_start = team_lists[i]
        team_link = team_lists[len(team_lists)-i-1]
        diff = get_diff(team_start, team_link)
        answer = min(answer, diff)

    print(answer)