DEV Community

Dev Cookies
Dev Cookies

Posted on

🚀 Day 8: Top K Elements Pattern (Amazon Interview Series)

Amazon loves problems where you must efficiently pick the "most" or "least" important items from a collection.

Instead of sorting everything (O(n log n)), we can use a heap/priority queue to maintain only the top K items in O(n log k).

🔑 Where Amazon Uses This Pattern

  • Finding Kth largest/smallest numbers
  • Getting Top K frequent words/numbers/products
  • Keeping K closest points to a location (Amazon Maps / Delivery optimization)
  • Streaming data processing (K largest seen so far)

📝 Problem 1: Kth Largest Element in an Array

👉 Amazon-style phrasing:
Given an integer array nums and an integer k, return the kth largest element.

Java Solution 

import java.util.*;

public class KthLargestElement {
    public static int findKthLargest(int[] nums, int k) {
        PriorityQueue<Integer> minHeap = new PriorityQueue<>();

        for (int num : nums) {
            minHeap.add(num);
            if (minHeap.size() > k) {
                minHeap.poll();
            }
        }
        return minHeap.peek();
    }

    public static void main(String[] args) {
        int[] nums = {3,2,1,5,6,4};
        System.out.println(findKthLargest(nums, 2)); // Output: 5
    }
}
Enter fullscreen mode Exit fullscreen mode

✅ Time: O(n log k)
✅ Space: O(k)

📝 Problem 2: Top K Frequent Elements

👉 Amazon-style phrasing:
Given an array of integers, 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> freq = new HashMap<>();
        for (int n : nums) freq.put(n, freq.getOrDefault(n, 0) + 1);

        PriorityQueue<Map.Entry<Integer, Integer>> minHeap =
            new PriorityQueue<>((a, b) -> a.getValue() - b.getValue());

        for (Map.Entry<Integer, Integer> entry : freq.entrySet()) {
            minHeap.add(entry);
            if (minHeap.size() > k) minHeap.poll();
        }

        int[] res = new int[k];
        for (int i = k - 1; i >= 0; i--) {
            res[i] = minHeap.poll().getKey();
        }
        return res;
    }

    public static void main(String[] args) {
        int[] nums = {1,1,1,2,2,3};
        System.out.println(Arrays.toString(topKFrequent(nums, 2))); 
        // Output: [1, 2]
    }
}
Enter fullscreen mode Exit fullscreen mode

✅ Amazon insight: Instead of sorting by frequency (O(n log n)), keep only top K in the heap.

📝 Problem 3: K Closest Points to Origin

👉 Amazon-style phrasing:
Given an array of points on a 2D plane, find the K closest points to the origin (0,0).

Java Solution 

public class KClosestPoints {
    public static int[][] kClosest(int[][] points, int K) {
        PriorityQueue<int[]> maxHeap = new PriorityQueue<>(
            (a, b) -> (b[0]*b[0] + b[1]*b[1]) - (a[0]*a[0] + a[1]*a[1])
        );

        for (int[] p : points) {
            maxHeap.add(p);
            if (maxHeap.size() > K) maxHeap.poll();
        }

        int[][] res = new int[K][2];
        for (int i = 0; i < K; i++) res[i] = maxHeap.poll();
        return res;
    }

    public static void main(String[] args) {
        int[][] points = {{3,3},{5,-1},{-2,4}};
        int[][] res = kClosest(points, 2);
        for (int[] p : res) System.out.println(Arrays.toString(p));
        // Output: [[3,3], [-2,4]]
    }
}
Enter fullscreen mode Exit fullscreen mode

✅ Why Amazon asks: It’s delivery-route / mapping-related — closest delivery stations, warehouses, etc.

📝 Problem 4: Connect Ropes (Minimum Cost)

👉 Amazon-style phrasing:
You’re given n ropes with different lengths. Connect them into one rope with minimum cost, where cost = sum of lengths joined each time.

Java Solution 

public class ConnectRopes {
    public static int minCost(int[] ropes) {
        PriorityQueue<Integer> minHeap = new PriorityQueue<>();
        for (int r : ropes) minHeap.add(r);

        int cost = 0;
        while (minHeap.size() > 1) {
            int a = minHeap.poll();
            int b = minHeap.poll();
            int newRope = a + b;
            cost += newRope;
            minHeap.add(newRope);
        }
        return cost;
    }

    public static void main(String[] args) {
        int[] ropes = {4,3,2,6};
        System.out.println(minCost(ropes)); // Output: 29
    }
}
Enter fullscreen mode Exit fullscreen mode

✅ Amazon insight: Warehouse or server merging optimization scenario.

📚 Extended Problem List (Amazon Loves These)

  1. Kth Smallest Element in a Sorted Matrix
  2. Top K Frequent Words (strings instead of numbers)
  3. Find Median from Data Stream (continuous input → heaps)
  4. Smallest Range Covering Elements from K Lists
  5. Sliding Window Median

🎯 Key Takeaways

  • Use Min Heap for Top K largest, Max Heap for Top K smallest.
  • Always compare O(n log k) vs naive sorting O(n log n).
  • Amazon expects you to discuss time-space tradeoffs during interviews.

📅 Next in the series (Day 9):
👉 Modified Binary Search Pattern – Amazon’s gold mine: rotated arrays, peak elements, search in unknown size arrays.

Top comments (0)