Amazon frequently asks problems where you need to extract K most/least elements.
π The two main tools are:
- Heap/PriorityQueue (straightforward, easy to reason)
- QuickSelect (optimized, avoids O(n log n))
π Problem 1: Kth Largest Element in an Array
π Amazon-style phrasing:
Find the kth largest element in an unsorted array.
β Java Solution (Min-Heap)
import java.util.*;
public class KthLargest {
public static int findKthLargest(int[] nums, int k) {
PriorityQueue<Integer> minHeap = new PriorityQueue<>();
for (int num : nums) {
minHeap.offer(num);
if (minHeap.size() > k) minHeap.poll();
}
return minHeap.peek();
}
}
β‘ Time Complexity: O(n log k)
β‘ Amazon Insight: Classic test of PriorityQueue usage.
π Problem 2: Top K Frequent Elements
π Amazon-style phrasing:
Given an integer array, return the k most frequent elements.
β Java Solution
import java.util.*;
public class TopKFrequent {
public static int[] topKFrequent(int[] nums, int k) {
Map<Integer, Integer> freqMap = new HashMap<>();
for (int n : nums) freqMap.put(n, freqMap.getOrDefault(n, 0) + 1);
PriorityQueue<Map.Entry<Integer, Integer>> heap =
new PriorityQueue<>((a, b) -> a.getValue() - b.getValue());
for (Map.Entry<Integer, Integer> entry : freqMap.entrySet()) {
heap.offer(entry);
if (heap.size() > k) heap.poll();
}
int[] result = new int[k];
for (int i = k - 1; i >= 0; i--) result[i] = heap.poll().getKey();
return result;
}
}
β‘ Time Complexity: O(n log k)
β‘ Amazon Insight: Very common β shows ability to combine HashMap + Heap.
π Problem 3: K Closest Numbers
π Amazon-style phrasing:
Given a sorted array and a number x, return the k closest numbers to x.
β Java Solution
import java.util.*;
public class KClosest {
public static List<Integer> findClosestElements(int[] arr, int k, int x) {
PriorityQueue<int[]> maxHeap =
new PriorityQueue<>((a, b) -> b[0] - a[0]);
for (int num : arr) {
maxHeap.offer(new int[]{Math.abs(num - x), num});
if (maxHeap.size() > k) maxHeap.poll();
}
List<Integer> result = new ArrayList<>();
while (!maxHeap.isEmpty()) result.add(maxHeap.poll()[1]);
Collections.sort(result);
return result;
}
}
β‘ Time Complexity: O(n log k)
β‘ Amazon Insight: They like this because it combines sorting + heap reasoning.
π Problem 4: Kth Smallest in a Sorted Matrix
π Amazon-style phrasing:
Given an n x n matrix sorted row & col-wise, find the kth smallest element.
β Java Solution
import java.util.*;
public class KthSmallestMatrix {
static class Node {
int val, row, col;
Node(int v, int r, int c) { val = v; row = r; col = c; }
}
public static int kthSmallest(int[][] matrix, int k) {
int n = matrix.length;
PriorityQueue<Node> minHeap =
new PriorityQueue<>((a, b) -> a.val - b.val);
for (int r = 0; r < Math.min(n, k); r++)
minHeap.offer(new Node(matrix[r][0], r, 0));
int count = 0, result = 0;
while (!minHeap.isEmpty()) {
Node node = minHeap.poll();
result = node.val;
if (++count == k) break;
if (node.col + 1 < n) {
minHeap.offer(new Node(matrix[node.row][node.col + 1],
node.row, node.col + 1));
}
}
return result;
}
}
β‘ Time Complexity: O(k log n)
β‘ Amazon Insight: Appears in medium/hard rounds β heap of pointers across rows.
π Problem 5: Kth Largest Element using QuickSelect (Optimized)
π When Amazon wants average O(n) solution, not heap.
β Java Solution
public class QuickSelectKthLargest {
public static int findKthLargest(int[] nums, int k) {
int target = nums.length - k;
int left = 0, right = nums.length - 1;
while (left <= right) {
int pivotIndex = partition(nums, left, right);
if (pivotIndex == target) return nums[pivotIndex];
else if (pivotIndex < target) left = pivotIndex + 1;
else right = pivotIndex - 1;
}
return -1;
}
private static int partition(int[] nums, int left, int right) {
int pivot = nums[right];
int i = left;
for (int j = left; j < right; j++) {
if (nums[j] <= pivot) {
swap(nums, i, j);
i++;
}
}
swap(nums, i, right);
return i;
}
private static void swap(int[] nums, int i, int j) {
int temp = nums[i]; nums[i] = nums[j]; nums[j] = temp;
}
}
β‘ Time Complexity: Average O(n), worst-case O(nΒ²).
β‘ Amazon Insight: Usually asked as a follow-up to βheap solutionβ.
π Extended Amazon Problem List (Heaps/Top-K)
- Merge K Sorted Lists (LeetCode 23)
- K Pairs with Smallest Sums
- Reorganize String (Greedy + Heap)
- Task Scheduler (Greedy + Heap)
- Sliding Window Maximum (Deque or Heap)
π― Key Takeaways
- Heap (PriorityQueue) β safest go-to for Top-K.
- QuickSelect β expected for optimization follow-up.
- Amazon likes HashMap + Heap combos (frequency, closest, etc.).
π
Next in the series (Day 13):
π Two Heaps Pattern β used in Median of Data Stream, Sliding Windows, etc.
Top comments (0)