[Python] 프로그래머스 - 빙고 (feat. set의 시간 복잡도)

문제 설명

N*N 숫자 입력과 지운 숫자가 들어가있는 nums 배열이 있다고 했을 때 총 몇 개의 빙고가 발생(?)하는지를 리턴하는 문제

 

나의 풀이

시간 초과 발생.

def solution(board, nums):
    answer = 0
    n = len(board)
    
    for row in range(n):
        for col in range(n):
            if board[row][col] in nums:
                board[row][col] = 0
    
    # 오른쪽으로 이동 (horizontal 빙고)
    for i in range(n):
        for j in range(n-1):
            if board[i][j] != board[i][j+1]:
                break
        else:
            answer += 1
            
    # 아래로 이동 (vertical 빙고)
    for i in range(n):
        for j in range(n-1):
            if board[j][i] != board[j+1][i]:
                break
        else:
            answer += 1
    
    # diagonal
    for i in range(n-1):
        if board[i][i] != board[i+1][i+1]:
            break
    else:
        answer += 1
           
    # off-diagonal
    for i in range(n-1):
        if board[i][n-1-i] != board[i+1][n-1-(i+1)]:
            break
    else:
        answer += 1
        
    return answer

로직 상으로는 문제가 없지만, 시간 초과가 떠서 효율성에서 0점 처리 됐다. 그 이유를 처음에는 horizontal, vertical 빙고를 찾으면서 $O(n^2)$의 시간복잡도를 갖기 때문이라고 생각했는데 빙고 처리를 아래와 같이 summation하는 방식으로 변경했는데도 시간 초과가 발생했다.

def solution(...):
	...
    
    # 각 빙고 방향에 대해서 다 룰 베이스로 판단하기.
    answer += len([i for i in board if sum(i) == 0]) # horizontal
    answer += len([i for i in zip(*board) if sum(i) == 0]) # vertical
    answer += int(sum(board[i][i] for i in range(n)) == 0) # diagonal
    answer += int(sum(board[n - 1 - i][i] for i in range(n)) == 0) # off-diagonal
    
    return answer

 

곰곰히 생각해보니, 문제는 빙고를 찾는 부분이 아니라 nums의 자리를 처리하는 부분이었다. 이미 for loop을 두 번 돌면서 $O(n^2)$의 시간복잡도가 발생했을 뿐 아니라 board의 특정 자리 값이 nums에 있는지를 확인하는 과정에서 $O(n)$의 시간복잡도가 추가적으로 들고, 결국엔 $O(n^3)$이라는 매우 큰 시간 복잡도를 갖게 되었다.

하지만, 이런 탐색 과정은 매우 필수적이기 때문에 탐색을 하지 않을 수는 없었다. 이에 좀 더 알아보니 nums는 리스트이기 때문에 탐색 시 $O(n)$의 시간 복잡도가 소요되지만, set 자료형을 사용하여 type-casting을 해주면 $O(1)$로 탐색이 가능하다. 이는 set이 해시테이블로 구현이 되었기 때문이다.

해시테이블 같은 경우에는 일반적으로 O(1)의 시간복잡도로 탐색이 가능하다. 해시에 대해서는 나중에 보다 자세히 정리해보도록 하겠다.

결국 set으로 nums를 타입 변경해줌으로써 결과적으로 bottleneck은 $O(n^2)$이 되었고, 이로써 시간 초과 문제는 해결하였다. 

 

로직 개선에 대한 추가 설명

def solution(...):
	...
    
    # 각 빙고 방향에 대해서 다 룰 베이스로 판단하기.
    answer += len([i for i in board if sum(i) == 0]) # horizontal
    answer += len([i for i in zip(*board) if sum(i) == 0]) # vertical
    answer += int(sum(board[i][i] for i in range(n)) == 0) # diagonal
    answer += int(sum(board[n - 1 - i][i] for i in range(n)) == 0) # off-diagonal
    
    return answer

vertical 빙고를 판단할 때에 zip함수에 board를 풀어줘서 넣어주는 방식이 매우 신기했다. 놀랍게도 저런 식으로 zip함수를 사용하면 matrix를 같은 열 & 다른 행에 있는 요소끼리 묶어줄 수 있었다.

 

임의로 (5,5)의 행렬을 만들어서 테스트해본 결과이다. print(*mat)을 수행했을 때, 각 행을 나열하였다. 따라서 이를 zip해주는 경우에는 각 행의 요소를 묶어서 리턴할 것이다. 따라서 다음과 같이 같은 열 & 다른 행의 요소를 묶어서 나열할 수 있었던 것이다. 위와 같이 행 간 summation을 이용해 vertical 빙고를 찾아내는 룰을 세울 수 있었다.

 

최종 코드

def solution(board, nums):
	answer = 0
    n = len(board)
    nums = set(nums)
    
    for row in range(n):
    	for col in range(n):
        	if board[row][col] in nums:
            	board[row][col] = 0

	# horizontal (각 행을 탐색하면서 sum = 0인지 체크)
    answer += len([i for i in board if sum(i)==0])
    # vertical (각 열을 탐색하면서 sum = 0인지 체크)
    answer += len([i for i in zip(*board) if sum(i)==0])
    # diagonal (대각요소의 합이 0인지 체크)
    answer += int(sum(board[i][i] for i in range(n))==0)
    # off-diagonal
    answer += int(sum(board[n-i-1][i] for i in range(n))==0)
    
    return answer

 

---

정리

• set은 해시테이블로 구현되어 있기 때문에 시간 복잡도가 $O(1)$이다.
• zip(*2D-array)를 수행하면, 같은 열 내 다른 행의 요소끼리 묶어줄 수 있다. (열 별로 slicing 가능)