Kth Smallest in Sorted Matrix

▶

Steps

setup

Initialize min-heap with first elements of each row

We push the first element of each row into the min-heap to start with the smallest candidates from each row.

💡 This step sets up the initial candidates for the kth smallest element, leveraging the sorted property of each row.

Line:

for r in range(min(k, n)):
    heapq.heappush(min_heap, (matrix[r][0], r, 0))

💡 The heap now contains the smallest element from each row, ready to be popped in ascending order.

compare

Pop smallest element from heap (1 from row 0, col 0)

We pop the smallest element 1 from the heap, which is the first candidate for the kth smallest element.

💡 Popping the smallest element helps us count how many elements have been processed in sorted order.

Line:

val, r, c = heapq.heappop(min_heap)
count += 1

💡 The smallest element is confirmed as the next in sorted order; count increments to 1.

insert

Push next element in row 0 (5 from row 0, col 1)

Since there is a next element in row 0, we push 5 into the heap to continue exploring that row.

💡 Pushing the next element maintains the heap with the next smallest candidate from the same row.

Line:

if c + 1 < n:
    heapq.heappush(min_heap, (matrix[r][c+1], r, c+1))

💡 The heap now contains the next smallest elements from all rows, including the newly pushed 5.

compare

Pop smallest element from heap (5 from row 0, col 1)

Pop 5, the smallest element currently in the heap, incrementing count to 2.

💡 Each pop confirms the next smallest element in sorted order.

Line:

val, r, c = heapq.heappop(min_heap)
count += 1

💡 The count tracks how many smallest elements have been found so far.

insert

Push next element in row 0 (9 from row 0, col 2)

Push 9 from row 0, column 2, as it is the next candidate in that row.

💡 Continuing to push next elements keeps the heap updated with smallest candidates.

Line:

if c + 1 < n:
    heapq.heappush(min_heap, (matrix[r][c+1], r, c+1))

💡 Heap now contains 9, 10, and 12, maintaining sorted order candidates.

compare

Pop smallest element from heap (9 from row 0, col 2)

Pop 9, increment count to 3.

💡 Popping continues to reveal the sorted order of elements.

Line:

val, r, c = heapq.heappop(min_heap)
count += 1

💡 We have now processed three smallest elements.

prune

No next element in row 0 to push

Row 0 is exhausted; no new element to push from this row.

💡 When a row is exhausted, the heap size decreases, focusing on remaining rows.

Line:

if c + 1 < n:
    heapq.heappush(min_heap, (matrix[r][c+1], r, c+1))

💡 Heap now only contains elements from rows 1 and 2.

compare

Pop smallest element from heap (10 from row 1, col 0)

Pop 10 from row 1, increment count to 4.

💡 We continue popping the smallest element to count towards kth smallest.

Line:

val, r, c = heapq.heappop(min_heap)
count += 1

💡 The heap now focuses on row 1 and 2 elements.

insert

Push next element in row 1 (11 from row 1, col 1)

Push 11 from row 1, column 1 into the heap.

💡 Pushing next element keeps the heap updated with smallest candidates.

Line:

if c + 1 < n:
    heapq.heappush(min_heap, (matrix[r][c+1], r, c+1))

💡 Heap now contains 11 and 12, continuing the merge.

compare

Pop smallest element from heap (11 from row 1, col 1)

Pop 11, increment count to 5.

💡 Each pop confirms the next smallest element in sorted order.

Line:

val, r, c = heapq.heappop(min_heap)
count += 1

💡 We are halfway to the 8th smallest element.

insert

Push next element in row 1 (13 from row 1, col 2)

Push 13 from row 1, column 2 into the heap.

💡 Continuing to push next elements keeps the heap updated.

Line:

if c + 1 < n:
    heapq.heappush(min_heap, (matrix[r][c+1], r, c+1))

💡 Heap now contains 12 and 13.

compare

Pop smallest element from heap (12 from row 2, col 0)

Pop 12, increment count to 6.

💡 Popping continues to reveal the sorted order.

Line:

val, r, c = heapq.heappop(min_heap)
count += 1

💡 Heap now contains 13.

insert

Push next element in row 2 (13 from row 2, col 1)

Push 13 from row 2, column 1 into the heap.

💡 Pushing next element keeps the heap updated with smallest candidates.

Line:

if c + 1 < n:
    heapq.heappush(min_heap, (matrix[r][c+1], r, c+1))

💡 Heap now contains two 13s from different rows.

compare

Pop smallest element from heap (13 from row 1, col 2)

Pop 13 from row 1, increment count to 7.

💡 Continuing to pop smallest elements to reach kth element.

Line:

val, r, c = heapq.heappop(min_heap)
count += 1

💡 Heap now contains one 13 from row 2.

prune

No next element in row 1 to push

Row 1 is exhausted; no new element to push.

💡 When a row is exhausted, the heap size decreases.

Line:

if c + 1 < n:
    heapq.heappush(min_heap, (matrix[r][c+1], r, c+1))

💡 Heap now contains only one element from row 2.

compare

Pop smallest element from heap (13 from row 2, col 1)

Pop 13 from row 2, increment count to 8, reaching the kth element.

💡 This pop returns the kth smallest element in the matrix.

Line:

val, r, c = heapq.heappop(min_heap)
count += 1
if count == k:
    return val

💡 The algorithm terminates with the correct kth smallest element.

from typing import List
import heapq

def kthSmallest(matrix: List[List[int]], k: int) -> int:
    n = len(matrix)
    min_heap = []
    # STEP 1: Initialize heap with first element of each row
    for r in range(min(k, n)):
        heapq.heappush(min_heap, (matrix[r][0], r, 0))
    count = 0
    while min_heap:
        # STEP 2: Pop smallest element
        val, r, c = heapq.heappop(min_heap)
        count += 1
        # STEP 3: Check if kth element reached
        if count == k:
            return val
        # STEP 4: Push next element in the same row if exists
        if c + 1 < n:
            heapq.heappush(min_heap, (matrix[r][c+1], r, c+1))

if __name__ == '__main__':
    matrix = [[1,5,9],[10,11,13],[12,13,15]]
    k = 8
    print(kthSmallest(matrix, k))  # Output: 13

📊

Kth Smallest in Sorted Matrix - Watch the Algorithm Execute, Step by Step

Watching each heap operation and element selection reveals how the algorithm narrows down the search space without sorting the entire matrix.

Target: 8

lo/mid

1

0

10

1

hi

12

2

mid[0]=1 - Initial smallest elements from each row

Next: lo=0 hi=2

Target: 8

lo/mid

10

0

hi

12

1

mid[0]=1 - Popped smallest element from heap

Next: lo=0 hi=1

Target: 8

lo

5

0

mid

10

1

hi

12

2

mid[1]=5 - Pushed next element from row 0

Next: lo=0 hi=2

Target: 8

lo/mid

10

0

hi

12

1

mid[0]=5 - Popped next smallest element

Next: lo=0 hi=1

Target: 8

lo/mid

9

0

10

1

hi

12

2

mid[0]=9 - Pushed next element from row 0

Next: lo=0 hi=2

Target: 8

lo/mid

10

0

hi

12

1

mid[0]=9 - Popped next smallest element

Next: lo=0 hi=1

Target: 8

lo

10

0

hi

12

1

Next: lo=0 hi=1

Target: 8

lo/mid/hi

12

0

mid[0]=10 - Popped next smallest element

Next: lo=0 hi=0

Target: 8

lo/mid

11

0

hi

12

1

mid[0]=11 - Pushed next element from row 1

Next: lo=0 hi=1

Target: 8

lo/mid/hi

12

0

mid[0]=11 - Popped next smallest element

Next: lo=0 hi=0

Target: 8

lo/mid

12

0

hi

13

1

mid[0]=13 - Pushed next element from row 1

Next: lo=0 hi=1

Target: 8

lo/mid/hi

13

0

mid[0]=12 - Popped next smallest element

Next: lo=0 hi=0

Target: 8

lo/mid

13

0

hi

13

1

mid[0]=13 - Pushed next element from row 2

Next: lo=0 hi=1

Target: 8

lo/mid/hi

13

0

mid[0]=13 - Popped next smallest element

Next: lo=0 hi=0

Target: 8

lo/hi

13

0

Next: lo=0 hi=0

Target: 8

Next: lo=0 hi=-1

Key Takeaways

✓ The min-heap efficiently merges sorted rows by always popping the smallest current candidate.

This insight is hard to see from code alone because the heap operations abstract away the merging process.

✓ Pushing the next element from the same row after popping maintains the heap with the next smallest candidates.

Visualizing this shows how the algorithm explores rows incrementally rather than sorting the entire matrix.

✓ Duplicates are naturally handled by the heap without special logic, as identical values appear and are popped in order.

Seeing duplicates processed step-by-step clarifies how the algorithm counts elements correctly.

Practice

1. Consider the following Python function that finds the kth positive missing number from a sorted array arr. What is the returned value when calling findKthPositive([2, 4, 7], 3)?

def findKthPositive(arr, k):
    n = len(arr)
    total_missing = arr[-1] - n
    if k > total_missing:
        return arr[-1] + (k - total_missing)
    left, right = 0, n - 1
    while left <= right:
        mid = (left + right) // 2
        missing = arr[mid] - (mid + 1)
        if missing < k:
            left = mid + 1
        else:
            right = mid - 1
    return arr[right] + (k - (arr[right] - (right + 1)))

easy

A. 5

B. 8

C. 9

D. 7

Solution

1. Step 1: Calculate total missing numbers

   Array length n=3, arr[-1]=7, total_missing = 7 - 3 = 4. Since k=3 ≤ 4, proceed with binary search.

2. Step 2: Perform binary search iterations

   Initial left=0, right=2. - mid=1, missing=arr[1] - (1+1) = 4 - 2 = 2 < 3 -> left=2 - mid=2, missing=arr[2] - (2+1) = 7 - 3 = 4 ≥ 3 -> right=1 Loop ends as left=2 > right=1. Return arr[right] + (k - (arr[right] - (right + 1))) = arr[1] + (3 - (4 - 2)) = 4 + (3 - 2) = 5.

3. Final Answer:

   Option A -> Option A

4. Quick Check:

   Binary search narrows to index 1, final missing number is 5 [OK]

Hint: Trace missing count at mid carefully [OK]

Common Mistakes:

• Off-by-one in missing calculation
• Returning arr[left] instead of arr[right]
• Miscomputing missing count as arr[mid] - mid

2. Consider the following Python code implementing an integer square root using Newton's method. What is the value of guess after the first iteration of the while loop when x = 8?

easy

A. 3

B. 4

C. 2

D. 1

Solution

1. Step 1: Initialize variables for x=8

   Initial guess = 8 // 2 = 4.

2. Step 2: First iteration update

   Check 4*4=16 > 8, so update guess = (4 + 8 // 4) // 2 = (4 + 2) // 2 = 3.

3. Final Answer:

   Option A -> Option A

4. Quick Check:

   First guess update correctly computes to 3 [OK]

Hint: Initial guess halves x; update uses integer division [OK]

Common Mistakes:

• Off-by-one in guess update
• Confusing integer division with float
• Miscomputing initial guess

3. Consider the following buggy feasibility check code for Capacity to Ship Packages. Which line contains the subtle bug that can cause incorrect feasibility results?

def can_ship_early_stop(weights, days, capacity):
    day_count = 1
    current_load = 0
    for w in weights:
        if current_load + w > capacity:
            day_count += 1
            current_load = 0
            if day_count > days:
                return False
        current_load += w
    return True

medium

A. Line checking if day_count exceeds days

B. Line incrementing day_count after resetting current_load

C. Line resetting current_load to 0 before incrementing day_count

D. Line adding current package weight to current_load

Solution

1. Step 1: Identify order of operations in feasibility check

   The bug is resetting current_load to 0 before incrementing day_count, which causes the current package weight to be lost.

2. Step 2: Correct approach

   Day count should increment first, then current_load reset to 0, then add current package weight to current_load.

3. Final Answer:

   Option C -> Option C

4. Quick Check:

   Resetting current_load before day_count causes lost package weight [OK]

Hint: Reset current_load after incrementing day_count and before adding weight [OK]

Common Mistakes:

• Swapping order of reset and increment
• Forgetting to increment day_count
• Adding weight before reset

4. The following code attempts to find the minimum in a rotated sorted array with duplicates. Identify the bug:

def find_min(nums):
    low, high = 0, len(nums) - 1
    while low < high:
        if nums[low] < nums[high]:
            return nums[low]
        mid = (low + high) / 2  # line 6
        if nums[mid] > nums[high]:
            low = mid + 1
        elif nums[mid] < nums[high]:
            high = mid
        else:
            high -= 1
    return nums[low]

What is the bug?

medium

A. Line 6 uses floating point division instead of integer division for mid calculation.

B. The condition nums[low] < nums[high] should be nums[low] <= nums[high].

C. The decrement of high in the else block should be increment of low instead.

D. Returning nums[low] at the end should be returning nums[mid].

Solution

1. Step 1: Examine mid calculation

   Line 6 uses (low + high) / 2 which results in a float in Python 3, causing an error when used as an index.

2. Step 2: Impact of bug

   Using float index causes runtime error or unexpected behavior; integer division (//) is required.

3. Final Answer:

   Option A -> Option A

4. Quick Check:

   Mid must be integer index [OK]

Hint: Mid index must be integer, not float [OK]

Common Mistakes:

• Using / instead of // for mid
• Returning nums[mid] instead of nums[low]

5. Consider the following buggy code snippet for checking if arrival is possible on time. Which line contains the subtle bug that can cause incorrect feasibility results?

from math import ceil

def can_arrive(dist, hour, speed):
    total_time = 0.0
    for i in range(len(dist)):
        total_time += ceil(dist[i] / speed)
    total_time -= ceil(dist[-1] / speed)  # Adjust last segment
    total_time += dist[-1] / speed
    return total_time <= hour

medium

A. Line 5: total_time += ceil(dist[i] / speed)

B. Line 6: total_time -= ceil(dist[-1] / speed) # Adjust last segment

C. Line 4: for i in range(len(dist))

D. Line 7: total_time += dist[-1] / speed

Solution

1. Step 1: Analyze loop bounds

   The loop iterates over all segments including the last, adding ceil for all, which overestimates time for last segment.

2. Step 2: Identify correct approach

   Only the first n-1 segments should use ceil; the last segment should add exact division time without ceiling.

3. Final Answer:

   Option C -> Option C

4. Quick Check:

   Loop should be for i in range(len(dist) - 1) [OK]

Hint: Last segment time must not be ceiled [OK]

Common Mistakes:

• Applying ceil to last segment
• Incorrect loop range including last segment