- The K Weakest Rows in a Matrix
링크텍스트
주어진 배열의 각각의 row에서 1의 개수가 가장 적은 순서대로 인덱스 k개를 출력하는 문제이다. 출력 시, i번째 row보다 j번째 row가 더 먼저 와야 한다. (i < j)
먼저 내가 풀었던 방법.
*lambda 사용법을 익히자
mat = [[1,1,0,0,0],
[1,1,1,1,0],
[1,0,0,0,0],
[1,1,0,0,0],
[1,1,1,1,1]]
k = 3
array = []
output = []
for i, row in enumerate(mat):
array.append([row.count(1), i])
# [[2, 0], [4, 1], [1, 2], [2, 3], [5, 4]]
array = sorted(array, key=lambda x:(x[0], x[1]))
# sorted 함수에서 key값은 정렬의 기준이다. 기본값은 오름차순.
# 값이 두 개 이상일 때, 조건을 넣어줄 수 있다.
# x[0], x[1]
# 배열의 첫번째 요소(x[0])를 기준으로 오름차순한 뒤
# 배열의 두번째 요소(x[1])를 기준으로 오름차순 한다.
# (첫번째 요소 기준으로 오름차순 정렬 했을때, 첫번째 요소의 값이 같은 애들끼리 오름차순 정렬한다.)
# [[1, 2], [2, 0], [2, 3], [4, 1], [5, 4]]
for i in range(k):
output.append(heap[i][1])
print(output)
heap을 사용한 간결한 코드
from heapq import heappush, heappop
mat = [[1,1,0,0,0],
[1,1,1,1,0],
[1,0,0,0,0],
[1,1,0,0,0],
[1,1,1,1,1]]
k = 3
heap = []
output = []
for i, row in enumerate(mat):
heappush(heap, [row.count(1), i]) # row 인덱스 '순서대로' push 하면 됨.
# [[1, 2], [2, 3], [2, 0], [4, 1], [5, 4]]
for i in range(k):
output.append(heappop(heap)[1])
# 이때 heap.pop()을 사용하면 [2, 3, 0] 순서대로 출력되니 주의하자.
# 꼭 heappop() 을 사용하여 힙의 구조가 유지되도록 해야한다. 올바른 출력: [2, 0, 3]
print(output)나는 heap을 사용할 때 자꾸 그냥 heap.pop()을 사용하게 되는데 꼭 heappop()을 사용하도록 주의하자.
