[CodeTree] 우리는 하나

 

https://www.codetree.ai/missions/2/problems/we-are-the-one/description

 

알고리즘

• 조합
• bfs

n x n개의 도시 중 k개의 도시를 선택한 뒤 각각의 경우에 대해 bfs를 진행하면 되는 이유다.

 

문제 자체는 평범하지만 시간 초과가 발생했고 조합에서 시간을 줄일 수 있는 문제라서 가져와보았다.

 

풀이 과정

변수

import sys
input = sys.stdin.readline
# 동서남북
dx = [0, 0, 1, -1]
dy = [1, -1, 0, 0]

# n : 격자의 크기 n x n
# k : 고를 도시의 수 
# u : 최저 높이 차
# d : 최고 높이 차
n, k, u, d = map(int, input().split())
arr = [list(map(int, input().split())) for _ in range(n)]

# 고른 좌표
selected_locations = []

 

기본적인 변수들이다.

 

좌표에서 동서남북 탐색을 위한 dx, dy를 정의했고 입력 값들을 정의했다.

 

selected_locations에는 선택한 도시의 좌표가 들어갈 예정이다.

 

bfs

# 갈 수 있는 서로 다른 도시의 수 
def get_different_cities(selected_locations):
    cnt = 0
    # 방문 여부
    visited = [[False]*n for _ in range(n)]
    for x, y in selected_locations:
        visited[x][y] = True
        cnt += 1

    # 선택된 도시들을 모두 큐에 넣음 
    queue = [location for location in selected_locations]
    while queue:
        x, y = queue.pop()
        for i in range(4):
            nx, ny = x + dx[i], y + dy[i]
            # 방문 가능하다면
            if 0<=nx<n and 0<=ny<n and not visited[nx][ny] and u<=abs(arr[nx][ny] - arr[x][y])<=d:
                visited[nx][ny] = True
                cnt+=1
                
                queue.append((nx, ny))
    return cnt

 

선택한 도시들을 바탕으로 갈 수 있는 서로 다른 도시의 개수를 구하는 함수이다.

 

선택한 도시들을 모두 큐에 넣고 bfs를 진행한다.

 

조합

max_res = k
# 현재까지 selected_cnt개를 뽑은 상태에서 [idx//n][idx%n] 좌표를 뽑을 것인지
def choose(idx, selected_cnt):
    if selected_cnt > k:
        return
    global max_res
    # 다 순회했다면 
    if idx == n**2:
        
        # 다 뽑았다면 
        if selected_cnt == k:
            visited_cnt = get_different_cities(selected_locations)
            max_res = max(max_res, visited_cnt)
        return
    # 뽑음
    selected_locations.append((idx // n, idx % n))
    choose(idx+1, selected_cnt + 1)
    selected_locations.pop()

    # 안 뽑음
    choose(idx+1, selected_cnt)
choose(0, 0)
print(max_res)

 

조합에서 시간을 단축할 수 있었던 핵심은 choose 함수 첫 부분에서 selected_cnt가 k가 넘어갔을 때 return하는 부분이다.

 

평소 조합 문제를 풀 때 이런 디테일까지 신경쓰지 않아도 시간 초과가 발생하지 않았는데 이 문제의 경우 이러한 처리가 시간을 많이 단축시켜 주었다.

 

문제를 보면 도시의 최대 개수는 8 x 8 = 64개이고 고를 도시의 최대 개수는 3이다.

 

즉 아무리 많아도 3개만 선택을 하는 문제인데 굳이 3개가 넘어가는 경우를 고려할 필요가 없다.

 

마무리

평소 아무 생각 없이 기계처럼 작성하던 조합 코드인데 문제에 맞춰 적절히 시간을 줄일 방법을 알게 되었다.

 

조합에 대한 이해가 부족했고 다음에 조합 문제에서 시간 초과가 발생한다면 이 방식을 고려해보려고 한다.