[Python/파이썬] 백준 17140번 이차원 배열과 연산

17140번: 이차원 배열과 연산

 

문제

크기가 3×3인 배열 A가 있다. 배열의 인덱스는 1부터 시작한다. 1초가 지날때마다 배열에 연산이 적용된다.

• R 연산: 배열 A의 모든 행에 대해서 정렬을 수행한다. 행의 개수 ≥ 열의 개수인 경우에 적용된다.
• C 연산: 배열 A의 모든 열에 대해서 정렬을 수행한다. 행의 개수 < 열의 개수인 경우에 적용된다.

한 행 또는 열에 있는 수를 정렬하려면, 각각의 수가 몇 번 나왔는지 알아야 한다. 그 다음, 수의 등장 횟수가 커지는 순으로, 그러한 것이 여러가지면 수가 커지는 순으로 정렬한다. 그 다음에는 배열 A에 정렬된 결과를 다시 넣어야 한다. 정렬된 결과를 배열에 넣을 때는, 수와 등장 횟수를 모두 넣으며, 순서는 수가 먼저이다.

예를 들어, [3, 1, 1]에는 3이 1번, 1가 2번 등장한다. 따라서, 정렬된 결과는 [3, 1, 1, 2]가 된다. 다시 이 배열에는 3이 1번, 1이 2번, 2가 1번 등장한다. 다시 정렬하면 [2, 1, 3, 1, 1, 2]가 된다.

정렬된 결과를 배열에 다시 넣으면 행 또는 열의 크기가 달라질 수 있다. R 연산이 적용된 경우에는 가장 큰 행을 기준으로 모든 행의 크기가 변하고, C 연산이 적용된 경우에는 가장 큰 열을 기준으로 모든 열의 크기가 변한다. 행 또는 열의 크기가 커진 곳에는 0이 채워진다. 수를 정렬할 때 0은 무시해야 한다. 예를 들어, [3, 2, 0, 0]을 정렬한 결과는 [3, 2]를 정렬한 결과와 같다.

행 또는 열의 크기가 100을 넘어가는 경우에는 처음 100개를 제외한 나머지는 버린다.

배열 A에 들어있는 수와 r, c, k가 주어졌을 때, A[r][c]에 들어있는 값이 k가 되기 위한 최소 시간을 구해보자.

 

입력

첫째 줄에 r, c, k가 주어진다. (1 ≤ r, c, k ≤ 100)

둘째 줄부터 3개의 줄에 배열 A에 들어있는 수가 주어진다. 배열 A에 들어있는 수는 100보다 작거나 같은 자연수이다.

출력

A[r][c]에 들어있는 값이 k가 되기 위한 연산의 최소 시간을 출력한다. 100초가 지나도 A[r][c] = k가 되지 않으면 -1을 출력한다.

코드

import heapq
from collections import defaultdict

r, c, k = map(int, input().split())
arr = [[0] * 100 for _ in range(100)]
for i in range(3):
    arr[i][0], arr[i][1], arr[i][2] = map(int, input().split())

# R 연산    
def row_sort():
    max_length = 0
    for i in range(100):
        num_cnt = defaultdict(int)
        for j in range(100):    # 숫자별 등장 횟수 카운트
            if arr[i][j] == 0:
                continue
            num_cnt[arr[i][j]] += 1
        
        # 수와 등장 횟수 정렬    
        heap = []
        for num, cnt in num_cnt.items():
            heapq.heappush(heap, (cnt, num))
            
        # 정렬된 결과 arr 배열에 넣기 
        idx = 0         
        while idx < 100:
            if heap:
                cnt, num = heapq.heappop(heap)
                arr[i][idx] = num
                arr[i][idx+1] = cnt
                idx += 2
                max_length = max(max_length, idx)
            else:
                arr[i][idx] = 0
                idx += 1
    return max_length
        
# C 연산    
def col_sort():
    max_length = 0
    for i in range(100):
        num_cnt = defaultdict(int)
        for j in range(100):    # 숫자별 등장 횟수 카운트
            if arr[j][i] == 0:
                continue
            num_cnt[arr[j][i]] += 1
            
        # 수와 등장 횟수 정렬    
        heap = []
        for num, cnt in num_cnt.items():
            heapq.heappush(heap, (cnt, num))
        
        # 정렬된 결과 arr 배열에 넣기
        idx = 0
        while idx < 100:
            if heap:
                cnt, num = heapq.heappop(heap)
                arr[idx][i] = num
                arr[idx+1][i] = cnt
                idx += 2
                max_length = max(max_length, idx)
            else:
                arr[idx][i] = 0
                idx += 1
    return max_length

row = 3
col = 3
for sec in range(101):
    if arr[r-1][c-1] == k:
        print(sec)
        exit()
    if row >= col:
        col = row_sort()
    else:
        row = col_sort()
    
print(-1)

 

처음에는 배열 arr에 필요한만큼의 숫자만 넣고 계산하려 했는데 문제를 풀다보니 어차피 모든 행의 크기를 맞춰줄 때 뒤에 0으로 채워주기도 하고 행 또는 열의 크기가 100을 넘어가는 경우에는 처음 100개를 제외한 나머지를 버린다고 했으니 0으로 초기화된 100 X 100 2차원 배열 arr을 만들어서 풀었다.

 

• R 연산 (row_sort())
  • 각 행의 숫자별 등장 횟수를 세서 defaultdict에 {숫자 : 등장횟수} 형태로 넣어준다. 만약 arr[i][j] == 0이라면 continue한다.
  • heapq를 이용하여 수와 등장횟수를 정렬한다. 이때 등장횟수가 커지는 순으로 먼저 정렬해야 하므로 이때는 (정렬횟수, 숫자) 형태로 heapq에 넣어준다.
  • heap에 저장된 값을 heappop하며 arr 배열에 넣어준다. 이때 heapq에서 pop해서 배열에 저장된 값들의 수가 열의 개수가 된다. 열의 개수를 알아야 나중에 R 연산을 할지 C 연산을 할지 결정할 수 있다. 따라서 이 값을 max_length와 비교하여 둘 중 큰 값을 max_length에 저장해준다. 이 max_length가 열의 크기이다.
    heap에 정렬된 값을 모두 pop했다면 그 뒤는 모두 0으로 업데이트해준다. 배열의 길이가 이전보다 짧아졌을 수도 있기에 뒤에는 모두 0으로 업데이트해주어야 한다.
• C 연산 (col_sort())
  • R 연산과 로직은 같다. max_length는 행의 크기가 된다.

 

 

문제를 다 풀고 다른 사람들은 어떻게 하였나 찾아봤더니 배열을 전치시켜서 풀이한 사람이 많았다.

배열을 전치시킨다는 건 행과 열을 바꾼다는 것인데 파이썬에서는 zip과 unpacking을 이용하여 코드 한 줄로 쉽게 배열을 전치시킬 수 있다.

mylist = [[1, 2, 3], [4, 5, 6], [7, 8, 9]]
new_list = list(map(list, zip(*mylist)))

이렇게 하면 

new_list는 [[1, 4, 7], [2, 5, 8], [3, 6, 9]]가 된다. 보기 쉽게 펼쳐보면 다음과 같다

[1, 4, 7]

[2, 5, 8]

[3, 6, 9]

파이썬에서는 이처럼 쉽게 배열 전치가 가능하다. 이를 이용하여 위의 문제를 풀었다면 굳이 행 연산과 열 연산을 위한 함수를 따로 만들지 않았어도 될 것이다.

 

[참고]

프로그래머스