K Weakest Squads
Problem
A training roster is a grid of 0s and 1s where each row is a squad. In every row the 1s (trained members) all come before the 0s (recruits). Squad i is weaker than squad j when it has fewer 1s, or the same number of 1s and a smaller index. Return the indices of the k weakest squads, weakest first.
Examples
Input: roster = [[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]
Why: the counts are 2, 4, 1, 2 and 5, and the tie between squads 0 and 3 goes to the smaller indexInput: roster = [[1, 0], [1, 0], [0, 0]], k = 2
Output: [2, 0]Input: roster = [[1, 1]], k = 1
Output: [0]
Why: edge case, a single squad is the weakest by defaultHints
0 / 3
Sorting every squad by strength works, but you only need k of them. What would you keep if the roster were too long to sort?
Keep the k weakest squads seen so far in a heap arranged so the strongest of them is on top. A new squad only matters if it is weaker than that top one, and then it replaces it.
Score each squad as (count of 1s, index). Push negated scores onto a min-heap so the strongest kept squad sits at the root, replace the root when a weaker squad arrives, and sort the k survivors weakest first at the end.
Solution
The k weakest squads are a top-k selection, so a heap of size k is enough: it holds the weakest squads seen so far with the strongest of them at the root, the one to evict next. Python's heapq is a min-heap, so each squad is stored as its negated (count, index) pair, which makes the strongest squad the smallest entry. A new squad replaces the root only when its pair is smaller than the root's, which is exactly the weaker-than rule including the index tie-break. Counting a row takes O(n) for n columns, so time is O(m · n + m log k) for m squads, and the heap uses O(k) space. Because each row is sorted, the count could also come from a binary search for the first 0.
import heapq
def k_weakest(roster, k):
heap = [] # k weakest so far, strongest at the root
for i, row in enumerate(roster):
entry = (-sum(row), -i) # negated: heapq keeps the smallest on top
if len(heap) < k:
heapq.heappush(heap, entry)
elif entry > heap[0]: # weaker than the strongest kept squad
heapq.heapreplace(heap, entry)
return [-i for _, i in sorted(heap, reverse=True)]
print(k_weakest([[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]], 3)) # -> [2, 0, 3]
print(k_weakest([[1, 0], [1, 0], [0, 0]], 2)) # -> [2, 0]
print(k_weakest([[1, 1]], 1)) # -> [0]Stuck on the idea rather than the code? Top K Elements covers it.