[알고리즘] 백준 스타트와 링크(14889번) Python

내가 한 풀이 : 조합

n의 범위가 최대 20이다.
따라서, O(n^3)에도 시간초과가 발생하지 않았다.

My solution

from itertools import combinations, permutations

# Calculate diff of teamA and teamB
def calculate(teamA, teamB):
    global stats
    sumA = 0
    sumB = 0
    memsA = permutations(teamA, 2)
    memsB = permutations(teamB, 2)

    for memA in memsA:
        sumA += stats[memA[0]][memA[1]]
    for memB in memsB:
        sumB += stats[memB[0]][memB[1]]

    result = abs(sumA - sumB)

    return result

# Initialize values
answer = 0
n = int(input())
stats = [[0 for j in range(n)] for i in range(n)]

# Input values
for i in range(n):
    stats[i] = list(map(int, input().split()))

# Deside members of team A, B for combination
teamA = list(combinations(range(n), int(n/2)))
for idx in range(len(teamA)):
    teamB = [k for k in range(n) if k not in teamA[idx]]
    tmp = calculate(teamA[idx], teamB)
    if idx == 0:
        answer = tmp
    elif tmp < answer:
        answer = tmp

print(answer)

보완할 점

• 조합을 통해 각팀의 멤버를 정하고, 또다시 순열을 통해 멤버들 중 2명으로 짝짓도록 했다.
• calculation 함수의 시간복잡도는 O(N^2)이다.
• 반면, 조금더 간단하게 구현할 수 있는 방법이 있다.
• teamA, B가 정해지면 바로 2중(2명씩 비교하니까)for문으로 값을 더하는거다.
• 물론 똑같이 시간복잡도는 O(N^2)이지만 코드가 간단해진다.
• 시간복잡도를 줄이기 위해서 dfs를 사용할 수 있다.

보완 코드

from itertools import combinations, permutations


# Initialize values
answer = 0
n = int(input())
stats = [[0 for j in range(n)] for i in range(n)]

# Input values
for i in range(n):
    stats[i] = list(map(int, input().split()))

# Deside members of team A, B for combination
teamA = list(combinations(range(n), int(n/2)))
for idx in range(len(teamA)):
    sumA = 0
    sumB = 0
    teamB = [k for k in range(n) if k not in teamA[idx]]
    # Find sumA, sumB and diff
    for x in teamA[idx]:
        for y in teamA[idx]:
            if x != y:
                sumA += stats[x][y]
    for x in teamB:
        for y in teamB:
            if x != y:
                sumB += stats[x][y]

    tmp = abs(sumA - sumB)

    # Find min value
    if idx == 0:
        answer = tmp
    elif tmp < answer:
        answer = tmp

print(answer)

오잉?! 길이는 짧아졌으나 시간이 더 걸린다....? 라이브러리가 아무래도 더 효율적인가보다 ㅠㅠ