[파이썬] 1337. The K Weakest Rows in a Matrix - 리트코드

문제

m x n이진 행렬 mat이 제공 되는데, 1의 개수가 가장 적은 k 개의 인덱스를 반환하는 프로그램을 작성하시오.

입력

Input: 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

출력

Output: [2,0,3]

작성한 코드1

def kWeakestRows(mat: List[List[int]], k: int) -> List[int]:
    ans = []
    s = collections.defaultdict(int)
    for i in range(len(mat)):
       # 1의 개수가 몇 개인지
        c = collections.Counter(mat[i])
        # 인덱스를 키로 1의 개수 만큼으로 하는 딕셔너리
        s[i] = c[1]
        
	# 값을 기준으로 정렬
    s = sorted(s.items(), key=lambda x:x[1])
    for key, val in s:
    	# key 값을 ans에 담는다.
        ans.append(key)

    return ans[:k]

작성한 코드2

def kWeakestRows2(mat: List[List[int]], k: int) -> List[int]:
    heap = []
    answer = []
    for i in range(len(mat)):
    	# i번째 리스트의 1의 갯수
        cnt = mat[i].count(1)
        #heap 리스트에 1의 갯수와 인덱스를 푸쉬
        heapq.heappush(heap, [cnt, i])
    for _ in range(k):
    	# pop하면 우선순위가 가장 작은 값부터
        answer.append(heapq.heappop(heap)[1])

    return answer

https://leetcode.com/problems/the-k-weakest-rows-in-a-matrix/submissions/

The K Weakest Rows in a Matrix - LeetCode