[Python] 백준 14889_스타트와 링크

2023. 2. 10. 11:02알고리즘

백트래킹 분류 문제다

DFS 탐색으로 경우의 수(조합)들을 구해서 풀면 된다


# 백준 14889_스타트와 링크
# N/2 명으로 이루어진 스타트팀과 링크팀이 축구함
# 두 팀의 능력치의 차이를 최소값 구하기
import sys

result = 10e9
def dfs(idx, score1):
    global result
    if len(a_team) == n // 2:
        score2 = 0
        b_team = [i for i in range(n) if i not in a_team]
        
        for i, b in enumerate(b_team):
            for j in range(i + 1, n//2):
                score2 += s[b][b_team[j]] + s[b_team[j]][b]
        
        result = min(result, abs(score1 - score2))

        return

    for i in range(idx + 1, n): # 0 ~ n - 1  까지의 사람과 모두
        temp = 0
        for j in a_team:
            temp += s[i][j] + s[j][i]  # a 팀에 속한 모든 사람들과의 능력치 계산

        a_team.append(i)
        dfs(i, score1 + temp)
        a_team.pop()

n = int(sys.stdin.readline().rstrip())
s = []  # s[i][j] 는 i번째 사람과 j번째 사람이 같은 팀일때 팀에 더해지는 능력치
a_team = []
for _ in range(n):
    s.append(list(map(int, sys.stdin.readline().rstrip().split())))

a_team.append(0)
dfs(0, 0)
print(result)

통과됐다!

그런데, A팀 B팀으로 배열을 나눠서 점수를 계산했는데, 소스가 비효율적인 것 같아서 좀더 고민해봤다.

A팀에 속하는 배열을 먼저 만든 다음에 A팀에 없는 애들을 B팀 배열에 추가하여 점수를 계산했는데 더 좋은 방법이 분명히 있을 것이다..!

그래서 생각해낸 것이 visited 배열

원래 DFS 할 때 visited 배열을 주로 쓰는데 왜 이번엔 그렇게 안썼는지는 의문이지만..


사실 A팀에 추가하기 위해 a_team 배열에 추가 삭제하는 append(), pop() 연산 자체는 둘다 O(1) 복잡도이며 visited 여부를 기록하기 위해서 visited[i] 로 원소에 접근하는 것도 O(1) 이라서 괜찮지만,

처음 소스에서는 dfs함수를 재귀 호출하기전 점수를 계산하는 부분이 있어 점수 계산 부분이 분산되어있어 보기 안좋다.

또, 팀원을 N/2 명으로 모두 나눴을때, B팀의 점수를 계산하기 위해서 n명을 다 탐색하면서 b_team 배열을 새로 만든 뒤 점수를 계산하기 때문에 소스가 복잡해진다.

 

하지만 visited 배열을 사용하면 visited[i]가 True인 애들은 A팀, False 인 애들을 B팀 점수로 계산하면 소스가 보기 좋아진다.

# 백준 14889_스타트와 링크
# N/2 명으로 이루어진 스타트팀과 링크팀이 축구함
# 두 팀의 능력치의 차이를 최소값 구하기
import sys

result = 10e9
def dfs(idx, count):
    global result
    if count == n // 2:
        score1, score2 = 0, 0

        for i in range(n - 1):
            for j in range(i + 1, n):
                if visited[i] == True and visited[j] == True:
                    score1 += s[i][j] + s[j][i]
                elif visited[i] == False and visited[j] == False:
                    score2 += s[i][j] + s[j][i]
                    
        result = min(result, abs(score1 - score2))

        return

    for i in range(idx, n): # 0 ~ n - 1  까지의 사람과 모두
        if not visited[i]:
            visited[i] = True
            dfs(i, count+1)
            visited[i] = False

n = int(sys.stdin.readline().rstrip())
s = []  # s[i][j] 는 i번째 사람과 j번째 사람이 같은 팀일때 팀에 더해지는 능력치
visited = [False] * n
for _ in range(n):
    s.append(list(map(int, sys.stdin.readline().rstrip().split())))

dfs(0, 0)
print(result)

근데 놀랍게도 시간은 거의 3~4초나 더 오래 걸림 ㅜㅜ

 

 

아무튼 이렇게 개선할 수 있었으니까 기록!

찝찝했던 부분 깔끔하게 수정할 수 있어서 좋았다. :)