코드

시간 초과 (Pypy3 통과)

n = int(input())
graph = []

for _ in range(n):
    graph.append(list(map(int, input().split())))
    
visited = [False for _ in range(n)]

# 두 팀 능력치의 최솟값을 담을 변수
diff = int(10e9)

def dfs(depth, idx):
    global diff

    # 팀이 2개 => 전체 사람의 절반만큼만 dfs를 수행해도 모든 선수에게 능력치를 부여할 수 있음
    if depth == (n//2):
        
        # [0] : 스타트 팀 능력치 누적 합
        # [1] : 링크 팀 능력치 누적 합
        stat = [0, 0]
        
        for i in range(n):
            for j in range(n):
                
                # i와 j가 같다 => 대각선 좌표 값 0
                # 굳이 연산할 필요 없음 (시간 절약)
                if i == j:
                    continue
                
                # 둘 다 방문한 곳은 스타트 팀 능력치
                if visited[i] and visited[j]:
                    stat[0] += graph[i][j]
                    
                # 둘 다 방문하지 않은 곳은 링크 팀 능력치
                elif not visited[i] and not visited[j]:
                    stat[1] += graph[i][j]
                    
        # 두 팀 능력치 최솟값 구하여 diff에 저장
        diff = min(diff, abs(stat[0] - stat[1]))
        return
    
    
    # 백트래킹 수행
    for i in range(idx, n):
        if not visited[i]:
            visited[i] = True
            dfs(depth+1, i+1)
            visited[i] = False
            
dfs(0, 0)
print(diff)

Python3 통과

from itertools import combinations

n = int(input())
graph = [list(map(int, input().split())) for _ in range(n)]

# 0~n-1까지 숫자가 들어간 set 생성
# graph의 좌표로 사용됨
all_team = set(i for i in range(n))

# 두팀 능력치의 차이
diff = int(10e9)

def dfs(start_team, idx):
    global diff
    
    # 입력받은 팀의 길이가 전체 인원의 절반이 되면 연산을 수행
    # => 팀은 2개고 n은 짝수이므로 인원이 절반이 되면
    # 최솟값을 구하고 재귀를 종료하면 된다
    if len(start_team) == (n//2):
        
        # set 끼리 - 연산을 사용하면 차집합이 나옴
        # 즉, 입력받은 팀에 소속되지 않은 팀은 다른 팀의 인원이 됨
        link_team = all_team - start_team
        
        # 차이의 최솟값을 구하는 연산
        diff = min(diff, abs(sum_stat(start_team) - sum_stat(link_team)))
        return
    
    # start_team에 add => i번째 사람이 등록됨
    # remove => i번째 사람을 뺌
    # ==> 백트래킹 수행
    for i in range(idx, n):
        start_team.add(i)
        dfs(start_team, i+1)
        start_team.remove(i)
        
# 능력치 계산 메소드
def sum_stat(team):
    stat = 0
    
    # 주어진 set을 combinations 메소드로 요소가 2개인 조합을 생성
    # 중복이 없이 생성되므로 이 값 자체가 graph의 좌표로 활용됨
    # 대각선을 기준으로 교차하는 좌표의 합이 해당 팀의 스탯이 되므로
    # 이 메소드에서 파라미터로 전달받은 team을 잘 나눠주면 각각 팀의 스탯을 구할 수 있음
    for a, b in combinations(team, 2):
        stat += graph[a][b] + graph[b][a]

    return stat

dfs(set(), 0)
print(diff)

문제 해설

백트래킹을 활용하여 푸는 문제.

방문한 좌표를 스타트 팀이라고 생각하면, 방문하지 않은 다른 좌표는 자동으로 링크 팀이 된다는 원리를 이용하여 문제를 해결하였다. 처음 Pypy3으로 통과한 코드는 시간 초과가 발생하였다. depth의 범위를 n의 절반까지만 연산하고 i와 j가 같은 경우는 반복에서 제외시켰음에도 시간 초과 문제가 해결되지 않아 구글링을 통해 시간 복잡도를 줄이는 방법을 찾았으며, 그 방법들은 아래와 같다.

  1. 해당 능력치 좌표에 방문한 로직을 visited 리스트가 아닌 set을 활용한다.
  2. itertools의 combinations 메소드를 활용하여 능력치들을 중복없이 생성한다.

위의 두 방법을 토대로 문제를 풀고 시간 초과를 벗어나 정답으로 인정받았다.

개인적으로 combinations는 파이썬에서 밖에 사용을 못하고, 코딩 테스트에서 사용하지 못하게 막는 경우도 있다고 하여 사용하지 않으려고 했으나 이 문제에서 쓰기에 너무 알맞아서 이 참에 확실하게 itertools에 대하여 공부할 수 있었다.

댓글남기기

참고