Advanced Data Structures: Heaps and Priority Queues in Java

Understanding and implementing Heaps (Min-Heap, Max-Heap) and Priority Queues in Java, and their applications in Dijkstra's algorithm and other scenarios.


Priority Queues in Java for Competitive Programming

Introduction to Priority Queues

A priority queue is an abstract data type that is like a regular queue or stack data structure, but has an additional feature. Each element in a priority queue is associated with a "priority". In a priority queue, elements are served (removed) based on their priority. Elements with higher priority are served before elements with lower priority, regardless of their order of insertion. If multiple elements have the same priority, they are typically served in the order they were enqueued (FIFO - First In, First Out), although this behavior can vary depending on the specific implementation.

Priority queues are incredibly useful in various algorithms and data structures, particularly in competitive programming where optimizing for speed and efficiency is crucial. Common applications include task scheduling, graph algorithms (like Dijkstra's and Prim's), and event simulation.

Understanding the Concept of Priority Queues and Heaps

Priority Queues

At a high level, you can imagine a priority queue as a line where people with urgent needs (high priority) cut to the front. This means that the next person served is not necessarily the one who has been waiting the longest, but the one with the most pressing situation. The core operations for a priority queue are:

  • `insert(element, priority)` or `enqueue(element, priority)`: Adds an element to the queue with a specific priority.
  • `peek()`: Returns the element with the highest priority without removing it.
  • `poll()` or `dequeue()`: Removes and returns the element with the highest priority.
  • `isEmpty()`: Checks if the queue is empty.
  • `size()`: Returns the number of elements in the queue.

Heaps and Priority Queues

A heap is a tree-based data structure that efficiently implements the priority queue abstract data type. It is a specialized tree-based data structure that satisfies the heap property. The heap property dictates the relationship between a parent node and its children.

There are two main types of heaps:

  • Min-Heap: The value of each node is less than or equal to the value of its children. The root node contains the smallest element.
  • Max-Heap: The value of each node is greater than or equal to the value of its children. The root node contains the largest element.

Heaps provide logarithmic time complexity (O(log n)) for insertion and deletion operations, making them an ideal choice for implementing priority queues. Most commonly, binary heaps are used due to their relatively simple implementation and good performance.

Therefore, the relationship is that heaps are a common and efficient implementation of priority queues. While other implementations are possible (e.g., using sorted arrays or linked lists), heaps generally provide the best balance of performance for the key operations.

Exploring Different Ways to Implement Priority Queues in Java

1. Using `java.util.PriorityQueue`

Java's built-in `PriorityQueue` class provides a heap-based implementation of a priority queue. By default, it's a min-heap (smallest element has the highest priority).

 import java.util.PriorityQueue;

public class PriorityQueueExample {
    public static void main(String[] args) {
        PriorityQueue<Integer> pq = new PriorityQueue<>(); // Min-Heap

        pq.offer(5);
        pq.offer(2);
        pq.offer(8);
        pq.offer(1);

        System.out.println("Peek: " + pq.peek()); // Output: 1

        while (!pq.isEmpty()) {
            System.out.println("Poll: " + pq.poll()); // Output: 1, 2, 5, 8
        }
    }
} 

2. Using `java.util.PriorityQueue` with a Custom Comparator

To create a max-heap or to prioritize elements based on a different criterion, you can use a custom `Comparator`.

 import java.util.PriorityQueue;
import java.util.Comparator;

public class PriorityQueueMaxHeapExample {
    public static void main(String[] args) {
        PriorityQueue<Integer> pq = new PriorityQueue<>(Comparator.reverseOrder()); // Max-Heap

        pq.offer(5);
        pq.offer(2);
        pq.offer(8);
        pq.offer(1);

        System.out.println("Peek: " + pq.peek()); // Output: 8

        while (!pq.isEmpty()) {
            System.out.println("Poll: " + pq.poll()); // Output: 8, 5, 2, 1
        }
    }
} 

3. Implementing a Priority Queue with a Custom Class and Comparator

This is useful when you need to store objects with multiple fields and prioritize based on a specific field.

 import java.util.PriorityQueue;
import java.util.Comparator;

class Task {
    String name;
    int priority;

    public Task(String name, int priority) {
        this.name = name;
        this.priority = priority;
    }

    @Override
    public String toString() {
        return "Task{" +
                "name='" + name + '\'' +
                ", priority=" + priority +
                '}';
    }
}

public class PriorityQueueCustomClassExample {
    public static void main(String[] args) {
        PriorityQueue<Task> pq = new PriorityQueue<>(Comparator.comparingInt(task -> task.priority)); // Min-Heap based on priority

        pq.offer(new Task("Low Priority Task", 3));
        pq.offer(new Task("High Priority Task", 1));
        pq.offer(new Task("Medium Priority Task", 2));

        while (!pq.isEmpty()) {
            System.out.println("Poll: " + pq.poll());
        }
        // Output:
        // Poll: Task{name='High Priority Task', priority=1}
        // Poll: Task{name='Medium Priority Task', priority=2}
        // Poll: Task{name='Low Priority Task', priority=3}
    }
} 

4. Implementing a Priority Queue using `TreeMap`

While less common, you could implement a priority queue using `TreeMap`. The keys in the `TreeMap` represent the priorities, and the values represent a list of elements with that priority. This approach provides sorted access based on priority. This is less efficient than a heap, but might be useful in certain niche cases where the priority values have specific properties.

 import java.util.TreeMap;
import java.util.ArrayList;
import java.util.List;

public class TreeMapPriorityQueue {

    private TreeMap<Integer, List<String>> map = new TreeMap<>();

    public void enqueue(String element, int priority) {
        if (!map.containsKey(priority)) {
            map.put(priority, new ArrayList<>());
        }
        map.get(priority).add(element);
    }

    public String dequeue() {
        if (map.isEmpty()) {
            return null; // Or throw an exception
        }
        int lowestPriority = map.firstKey();
        List<String> elements = map.get(lowestPriority);
        String element = elements.remove(0);
        if (elements.isEmpty()) {
            map.remove(lowestPriority);
        }
        return element;
    }

    public String peek() {
        if (map.isEmpty()) {
            return null; // Or throw an exception
        }
        int lowestPriority = map.firstKey();
        List<String> elements = map.get(lowestPriority);
        return elements.get(0);
    }


    public boolean isEmpty() {
        return map.isEmpty();
    }

    public static void main(String[] args) {
        TreeMapPriorityQueue pq = new TreeMapPriorityQueue();
        pq.enqueue("Task C", 3);
        pq.enqueue("Task A", 1);
        pq.enqueue("Task B", 2);
        pq.enqueue("Task D", 1);


        System.out.println("Peek: " + pq.peek()); // Task A
        System.out.println("Dequeue: " + pq.dequeue()); //Task A
        System.out.println("Dequeue: " + pq.dequeue()); //Task D
        System.out.println("Dequeue: " + pq.dequeue()); //Task B
        System.out.println("Dequeue: " + pq.dequeue()); //Task C
        System.out.println("isEmpty: " + pq.isEmpty()); //true


    }
} 

Priority Queues in Competitive Programming with Solutions

Problem 1: K Largest Elements

Problem Statement: Given an array of integers, find the k largest elements.

Solution Approach: Use a min-heap (PriorityQueue) of size k. Iterate through the array. If the current element is larger than the smallest element in the heap (heap's root), remove the smallest element and insert the current element. After iterating through the entire array, the heap will contain the k largest elements.

 import java.util.PriorityQueue;
import java.util.Arrays;

public class KLargestElements {
    public static int[] findKLargest(int[] nums, int k) {
        PriorityQueue<Integer> minHeap = new PriorityQueue<>(); // Min-Heap

        for (int num : nums) {
            if (minHeap.size() < k) {
                minHeap.offer(num);
            } else if (num > minHeap.peek()) {
                minHeap.poll();
                minHeap.offer(num);
            }
        }

        int[] result = new int[k];
        for (int i = 0; i < k; i++) {
            result[i] = minHeap.poll();
        }

        Arrays.sort(result); //optional, sort for consistent results.
        return result;
    }

    public static void main(String[] args) {
        int[] nums = {3, 1, 4, 1, 5, 9, 2, 6, 5, 3, 5};
        int k = 3;
        int[] kLargest = findKLargest(nums, k);
        System.out.println("K Largest Elements: " + Arrays.toString(kLargest)); // Output: K Largest Elements: [5, 6, 9]
    }
} 
Explanation: The min-heap maintains the `k` smallest elements seen so far. If the current number is bigger than the smallest in the heap, then it must be amongst the k-largest. Thus the smallest is replaced with the current element.

Problem 2: Merge K Sorted Lists

Problem Statement: Given k sorted linked lists, merge them into one sorted linked list.

Solution Approach: Use a min-heap (PriorityQueue) to store the heads of all k lists. The heap will be ordered based on the value of the list head. Repeatedly remove the smallest element from the heap (which is the head of one of the lists), add it to the merged list, and then add the next element from the same list back into the heap (if the list isn't empty).

 import java.util.PriorityQueue;

class ListNode {
    int val;
    ListNode next;
    ListNode(int val) { this.val = val; }
}

public class MergeKSortedLists {

    public ListNode mergeKLists(ListNode[] lists) {
        PriorityQueue<ListNode> pq = new PriorityQueue<>((a, b) -> a.val - b.val); // Min-Heap based on ListNode's val

        // Add the head of each list to the priority queue
        for (ListNode node : lists) {
            if (node != null) {
                pq.offer(node);
            }
        }

        ListNode dummy = new ListNode(0); // Dummy node to build the merged list
        ListNode tail = dummy;

        while (!pq.isEmpty()) {
            ListNode node = pq.poll();
            tail.next = node;
            tail = node;

            if (node.next != null) {
                pq.offer(node.next);
            }
        }

        return dummy.next;
    }

    public static void main(String[] args) {
        // Example usage (create some sorted linked lists)
        ListNode list1 = new ListNode(1);
        list1.next = new ListNode(4);
        list1.next.next = new ListNode(5);

        ListNode list2 = new ListNode(1);
        list2.next = new ListNode(3);
        list2.next.next = new ListNode(4);

        ListNode list3 = new ListNode(2);
        list3.next = new ListNode(6);

        ListNode[] lists = {list1, list2, list3};

        MergeKSortedLists merger = new MergeKSortedLists();
        ListNode mergedList = merger.mergeKLists(lists);

        // Print the merged list
        while (mergedList != null) {
            System.out.print(mergedList.val + " ");
            mergedList = mergedList.next;
        }
        // Output: 1 1 2 3 4 4 5 6
    }
} 
Explanation: The priority queue holds the *heads* of all the lists. We repeatedly take the smallest head and add it to the resulting list. After extracting a head, we add the *next* node in that extracted head's list to the queue, so that it becomes a candidate to be the next smallest.

Problem 3: Dijkstra's Algorithm for Shortest Path

Problem Statement: Find the shortest path from a source node to all other nodes in a weighted graph.

Solution Approach: Dijkstra's algorithm uses a priority queue to efficiently select the node with the smallest distance from the source that hasn't been visited yet. The priority queue is ordered by the distance from the source node.

 import java.util.*;

public class Dijkstra {

    public static int[] dijkstra(int[][] graph, int source) {
        int n = graph.length;
        int[] dist = new int[n];
        Arrays.fill(dist, Integer.MAX_VALUE);
        dist[source] = 0;

        PriorityQueue<Integer> pq = new PriorityQueue<>(Comparator.comparingInt(a -> dist[a])); //Min-Heap based on distance from source.
        pq.offer(source);

        while (!pq.isEmpty()) {
            int u = pq.poll();

            for (int v = 0; v < n; v++) {
                if (graph[u][v] != 0) { // There is an edge between u and v
                    if (dist[u] != Integer.MAX_VALUE && dist[u] + graph[u][v] < dist[v]) {
                        dist[v] = dist[u] + graph[u][v];
                        pq.offer(v);  // Re-add to PQ if a shorter path is found.
                    }
                }
            }
        }

        return dist;
    }

    public static void main(String[] args) {
        int[][] graph = {
                {0, 4, 0, 0, 0, 0, 0, 8, 0},
                {4, 0, 8, 0, 0, 0, 0, 11, 0},
                {0, 8, 0, 7, 0, 4, 0, 0, 2},
                {0, 0, 7, 0, 9, 14, 0, 0, 0},
                {0, 0, 0, 9, 0, 10, 0, 0, 0},
                {0, 0, 4, 14, 10, 0, 2, 0, 0},
                {0, 0, 0, 0, 0, 2, 0, 1, 6},
                {8, 11, 0, 0, 0, 0, 1, 0, 7},
                {0, 0, 2, 0, 0, 0, 6, 7, 0}
        };

        int source = 0;
        int[] shortestDistances = dijkstra(graph, source);

        System.out.println("Shortest distances from node " + source + ":");
        for (int i = 0; i < shortestDistances.length; i++) {
            System.out.println("To node " + i + ": " + shortestDistances[i]);
        }
    }
} 
Explanation: The priority queue ensures that we always explore the node that has the *shortest known distance from the source*. By repeatedly visiting the node with the shortest distance and updating the distances to its neighbors, we gradually discover the shortest path to all other nodes. Note: Re-adding to the PQ is important for correctness if a shorter path is discovered *after* a node has already been visited. A more efficient solution would be to update the value associated with the key in the priority queue. This requires using a data structure that allows efficiently updating keys, such as a Fibonacci heap or a balanced binary search tree.