DEV Community

Dev Cookies
Dev Cookies

Posted on

πŸ“Š Day 12: Top-K Elements – Amazon Interview Series

Amazon frequently asks problems where you need to extract K most/least elements.
πŸ‘‰ The two main tools are:

  1. Heap/PriorityQueue (straightforward, easy to reason)
  2. 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();
    }
}
Enter fullscreen mode Exit fullscreen mode

⚑ 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;
    }
}
Enter fullscreen mode Exit fullscreen mode

⚑ 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;
    }
}
Enter fullscreen mode Exit fullscreen mode

⚑ 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;
    }
}
Enter fullscreen mode Exit fullscreen mode

⚑ 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;
    }
}
Enter fullscreen mode Exit fullscreen mode

⚑ 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)

  1. Merge K Sorted Lists (LeetCode 23)
  2. K Pairs with Smallest Sums
  3. Reorganize String (Greedy + Heap)
  4. Task Scheduler (Greedy + Heap)
  5. 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)