Kth Largest Element in Array

▶

Steps

setup

Start Quickselect

Initialize the Quickselect call with the full array and k adjusted to zero-based index (k-1 = 1).

💡 Setting up the initial parameters defines the search boundaries for the kth largest element.

Line:

return quickselect(nums, 0, len(nums) - 1, k - 1)

💡 The algorithm will search between indices 0 and 5 for the element at index 1 in the sorted order.

partition

Partition: Choose pivot

Select the pivot as the rightmost element, 4, to partition the array around it.

💡 Choosing the pivot is crucial as it determines how the array will be split into greater and smaller elements.

Line:

pivot = nums[right]

💡 Pivot 4 will be used to rearrange elements so that those greater than 4 come before it.

compare

Partition: Compare nums[0] with pivot

Compare element 3 at index 0 with pivot 4; since 3 is not greater than 4, no swap occurs.

💡 Comparisons determine if elements should be moved to the left side of the pivot.

Line:

if nums[j] > pivot:

💡 Elements not greater than pivot remain in place, preserving partition order.

compare

Partition: Compare nums[1] with pivot

Compare element 2 at index 1 with pivot 4; 2 is not greater than 4, so no swap.

💡 Each element is checked to decide if it should be moved left of the pivot.

Line:

if nums[j] > pivot:

💡 Elements smaller or equal to pivot remain on the right side after partition.

compare

Partition: Compare nums[2] with pivot

Compare element 1 at index 2 with pivot 4; 1 is not greater than 4, no swap.

💡 Continue checking elements to find those greater than pivot.

Line:

if nums[j] > pivot:

💡 Elements smaller than pivot accumulate on the right side.

compare

Partition: Compare nums[3] with pivot

Compare element 5 at index 3 with pivot 4; 5 is greater, so swap with nums[i] at index 0.

💡 Swapping moves elements greater than pivot to the left side.

Line:

if nums[j] > pivot:
    nums[i], nums[j] = nums[j], nums[i]
    i += 1

💡 Partitioning rearranges array to group larger elements left of pivot.

compare

Partition: Compare nums[4] with pivot

Compare element 6 at index 4 with pivot 4; 6 is greater, swap with nums[i] at index 1.

💡 Continue moving greater elements to the left side of the pivot.

Line:

if nums[j] > pivot:
    nums[i], nums[j] = nums[j], nums[i]
    i += 1

💡 Partitioning groups all elements greater than pivot at the front.

partition

Partition: Final swap pivot into place

Swap pivot 4 with element at index i=2 to place pivot in correct position.

💡 Placing pivot finalizes partition, separating greater and smaller elements.

Line:

nums[i], nums[right] = nums[right], nums[i]

💡 Pivot is now at index 2, with greater elements left and smaller right.

compare

Check pivot position against k

Pivot index 2 is greater than target k=1, so recurse on left subarray [0..1].

💡 Deciding which side to recurse on narrows the search space efficiently.

Line:

if pivot_index == k:
    return nums[pivot_index]
elif pivot_index < k:
    return quickselect(nums, pivot_index + 1, right, k)
else:
    return quickselect(nums, left, pivot_index - 1, k)

💡 The kth largest element must be in the left partition.

partition

Partition left subarray: Choose pivot

Choose pivot 6 at index 1 for left subarray [5,6].

💡 Pivot choice in the smaller subarray continues partitioning.

Line:

pivot = nums[right]

💡 Pivot 6 will separate elements greater than 6 (none) from smaller.

compare

Partition left subarray: Compare nums[0] with pivot

Compare 5 at index 0 with pivot 6; 5 is not greater, no swap.

💡 Check if elements are greater than pivot to move left.

Line:

if nums[j] > pivot:

💡 Elements smaller than pivot remain right of it.

partition

Partition left subarray: Final swap pivot into place

Swap pivot 6 with element at index i=0 to place pivot correctly.

💡 Pivot placement finalizes partition in the subarray.

Line:

nums[i], nums[right] = nums[right], nums[i]

💡 Pivot 6 is now at index 0, with greater elements left (none) and smaller right (5).

compare

Check pivot position against k

Pivot index 0 is less than k=1, so recurse on right subarray [1..1].

💡 Narrowing search to the right side where kth largest lies.

Line:

if pivot_index == k:
    return nums[pivot_index]
elif pivot_index < k:
    return quickselect(nums, pivot_index + 1, right, k)
else:
    return quickselect(nums, left, pivot_index - 1, k)

💡 The kth largest element is in the right partition of the subarray.

setup

Base case reached

Left and right pointers both at index 1; return element 5 as kth largest.

💡 When search space is one element, that element is the answer.

Line:

if left == right:
    return nums[left]

💡 The algorithm found the kth largest element without full sorting.

import random

def partition(nums, left, right):
    pivot = nums[right]  # STEP 2, 10
    i = left
    for j in range(left, right):
        if nums[j] > pivot:  # STEP 3-7, 11
            nums[i], nums[j] = nums[j], nums[i]  # STEP 6,7,12
            i += 1
    nums[i], nums[right] = nums[right], nums[i]  # STEP 8,12
    return i

def quickselect(nums, left, right, k):
    if left == right:  # STEP 14
        return nums[left]
    pivot_index = partition(nums, left, right)  # STEP 2-8,10-12
    if pivot_index == k:  # STEP 9,13
        return nums[pivot_index]
    elif pivot_index < k:
        return quickselect(nums, pivot_index + 1, right, k)
    else:
        return quickselect(nums, left, pivot_index - 1, k)

def findKthLargest(nums, k):
    return quickselect(nums, 0, len(nums) - 1, k - 1)  # STEP 1

if __name__ == '__main__':
    print(findKthLargest([3,2,1,5,6,4], 2))  # Output: 5

📊

Kth Largest Element in Array - Watch the Algorithm Execute, Step by Step

Watching each partition and pivot decision helps you understand how Quickselect efficiently finds the kth largest element without fully sorting the array.

setupmax-heap

3

0

2

1

1

2

5

3

6

4

4

5

partitionmax-heap

3

0

2

1

1

2

5

3

6

4

4

5

partitionmax-heap

3

0

2

1

1

2

5

3

6

4

4

5

partitionmax-heap

3

0

2

1

1

2

5

3

6

4

4

5

partitionmax-heap

3

0

2

1

1

2

5

3

6

4

4

5

partitionmax-heapswap [0↔3]

5

0

2

1

1

2

3

3

6

4

4

5

partitionmax-heapswap [1↔4]

5

0

6

1

1

2

3

3

2

4

4

5

partitionmax-heapswap [2↔5]

5

0

6

1

4

2

3

3

2

4

1

5

Added: 2

traversemax-heap

5

0

6

1

4

2

3

3

2

4

1

5

partitionmax-heap

5

0

6

1

4

2

3

3

2

4

1

5

partitionmax-heap

5

0

6

1

4

2

3

3

2

4

1

5

partitionmax-heapswap [0↔1]

6

0

5

1

4

2

3

3

2

4

1

5

Added: 0

traversemax-heap

6

0

5

1

4

2

3

3

2

4

1

5

peekmax-heap

6

0

5

1

4

2

3

3

2

4

1

5

Added: 5

Key Takeaways

✓ Quickselect efficiently finds the kth largest element by partitioning rather than fully sorting.

Reading code alone hides how partitioning narrows the search space step-by-step.

✓ Pivot selection and partitioning rearrange elements so that the kth largest element is isolated.

Visualizing swaps and comparisons clarifies how elements move relative to the pivot.

✓ The recursive decision to search left or right subarray depends on pivot position relative to k.

Seeing the pivot index compared to k helps understand the algorithm's divide-and-conquer nature.

Practice

1. Consider the following Python code implementing the heap-based approach to find the k-th smallest pair distance. What is the returned value when calling smallestDistancePair([1,3,1], 2)?

easy

A. 3

B. 1

C. 0

D. 2

Solution

1. Step 1: Sort input and initialize heap

   Sorted nums = [1,1,3]. Initial heap contains pairs: (0,0,1), (2,1,2).

2. Step 2: Pop k-1=1 times and push next pairs

   Pop (0,0,1). Since j+1=2<n, push (3-1=2,0,2). Heap now: (2,0,2), (2,1,2). The smallest distance at heap top is 2.

3. Step 3: Return the k-th smallest distance

   After popping once, the top of the heap is (2,0,2), so the returned value is 2.

4. Final Answer:

   Option D -> Option D

5. Quick Check:

   Second smallest pair distance is 2 [OK]

Hint: Trace heap pops and pushes carefully for small inputs [OK]

Common Mistakes:

• Off-by-one in k-1 loop
• Confusing indices i,j
• Returning wrong heap element

2. Consider the following Python code implementing the optimal approach for top K frequent words. Given the input words = ["apple", "banana", "apple", "orange", "banana", "banana"] and k = 2, what is the returned list?

easy

A. ["banana", "apple"]

B. ["banana", "orange"]

C. ["apple", "banana"]

D. ["apple", "orange"]

Solution

1. Step 1: Count frequencies

   apple: 2, banana: 3, orange: 1; max_freq = 3

2. Step 2: Fill buckets and sort within buckets

   buckets[3] = ["banana"], buckets[2] = ["apple"], buckets[1] = ["orange"]; iterate from freq=3 down

3. Final Answer:

   Option A -> Option A

4. Quick Check:

   Top 2 frequent words are banana (3), apple (2) [OK]

Hint: Trace frequency buckets from highest to lowest [OK]

Common Mistakes:

• Confusing lex order with frequency order
• Off-by-one in bucket indices
• Returning words in ascending frequency order

3. What is the time complexity of the optimal algorithm that uses a prefix XOR matrix and a min-heap of size k to find the kth largest XOR coordinate value in an m x n matrix?

medium

A. O(m * n * k)

B. O(m * n * log k)

C. O(m^2 * n^2)

D. O(m * n + k * log(m * n))

Solution

1. Step 1: Identify loops and heap operations

   The algorithm iterates over all m*n coordinates, computing prefix XOR in O(1) per cell.

2. Step 2: Analyze heap operations per coordinate

   For each coordinate, heap operations (push or replace) take O(log k) time, since heap size is limited to k.

3. Final Answer:

   Option B -> Option B

4. Quick Check:

   Overall complexity is O(m * n * log k), not O(m * n * k) or O(m^2 * n^2) [OK]

Hint: Heap operations cost log k, not linear in k [OK]

Common Mistakes:

• Confusing k with n or m
• Assuming sorting all values is needed
• Ignoring heap operation cost

4. Suppose the frequency stack must support a new operation popK(k) that removes and returns the top k most frequent elements in order, breaking ties by recency. Which modification to the optimal heap-based approach correctly supports this operation efficiently?

hard

A. Perform k consecutive pop() calls on the heap, updating frequencies after each pop.

B. Use a separate max-heap for each frequency level and merge results on popK(k).

C. Modify the heap to store lists of elements per frequency and pop from these lists in order.

D. Maintain a global sorted list of elements by frequency and recency, and slice the top k on popK(k).

Solution

1. Step 1: Understand popK(k) requirements

   We must pop the top k elements by frequency and recency order, updating frequencies after each pop.

2. Step 2: Evaluate modifications

   Perform k consecutive pop() calls on the heap, updating frequencies after each pop. This performs k pops sequentially, each O(log n), correctly updating frequencies and maintaining heap invariants. Options A, B, and D either complicate the data structure or lose efficiency and correctness.

3. Final Answer:

   Option A -> Option A

4. Quick Check:

   Sequential pops maintain correctness and efficiency [OK]

Hint: Sequential pops maintain heap correctness for top-k [OK]

Common Mistakes:

• Trying to pop multiple elements at once without updating frequencies
• Using complex data structures that break heap invariants

5. Suppose the problem changes: you must find the k closest points to the origin, but points can be repeated multiple times (i.e., duplicates allowed and must be counted separately). Which modification to the Quickselect approach correctly handles this scenario without increasing time complexity significantly?

hard

A. Use a max heap of size k to keep track of closest points, allowing duplicates naturally.

B. Use a min heap to store all points and pop k times to get closest points, handling duplicates automatically.

C. Sort all points by distance and then remove duplicates before returning the first k points.

D. Modify Quickselect to partition points by distance and include duplicates in the partitioning logic without skipping any points.

Solution

1. Step 1: Understand duplicates impact

   Duplicates must be treated as distinct points; Quickselect partitions by distance and does not rely on uniqueness.

2. Step 2: Modify partition to handle duplicates correctly

   Partition logic remains the same; all points, including duplicates, are compared and placed correctly without skipping.

3. Final Answer:

   Option D -> Option D

4. Quick Check:

   Quickselect handles duplicates naturally if partitioning is correct [OK]

Hint: Quickselect partitions all points regardless of duplicates [OK]

Common Mistakes:

• Using min heap instead of max heap
• Removing duplicates before selection
• Sorting all points unnecessarily