Search code examples
javapriority-queuemin-heap

Use PriorityQueue as a minHeap


When trying to solve this problem https://www.hackerrank.com/challenges/cut-the-tree I was trying to always pick and cut a leaf and then combine its weight to the node that it connects. I was using a PriorityQueue to store all the nodes, and using the size of their adjacent nodes as the priority. But when I'm trying some test case, it seems that the priority queue property is violated, which means that non-leaf nodes may appear before leaf nodes. Will PriorityQueue automatically update itself or should I call some function to update it. My temporary solution is to use a list to store all the leaves.

The following is my code:

public class Solution {
    private static class Node implements Comparable<Node> {
        int index;
        int value;
        Map<Integer, Node> adj;

        public Node(int index, int value) {
            this.index = index;
            this.value = value;
            this.adj = new HashMap<Integer, Node>();
        }

        @Override
        public int compareTo(Node n) {
            return adj.size() - n.adj.size();
        }
    }

    public static void main(String[] args) {

        BufferedReader br = null;
        try {
            br = new BufferedReader(new InputStreamReader(new FileInputStream("input.txt")));

            int total = 0;

            int N = Integer.parseInt(br.readLine());
            String[] strs = br.readLine().split(" ");
            HashMap<Integer, Node> nodes = new HashMap<Integer, Node>();

            for (int i = 0; i < N; i++) {
                int value = Integer.parseInt(strs[i]);
                nodes.put(i, new Node(i, value));
                total += value;
            }

            for (int i = 0; i < N - 1; i++) {
                strs = br.readLine().split(" ");
                int n1 = Integer.parseInt(strs[0]) - 1;
                int n2 = Integer.parseInt(strs[1]) - 1;

                nodes.get(n1).adj.put(n2, nodes.get(n2));
                nodes.get(n2).adj.put(n1, nodes.get(n1));
            }

            // PriorityQueue<Node> pq = new PriorityQueue<Node>(nodes.values());
            // while (!pq.isEmpty()) {
            // Node n = pq.poll();
            // System.out.println(n.index + " " + n.adj.size());
            // }
            // NOTE: java's PriorityQueue doesn't support update, cannot use it
            // LAME DESIGN. use a LinkedList instead

            List<Node> leaves = new LinkedList<Node>();
            for (Node node : nodes.values()) {
                if (node.adj.size() == 1) {
                    leaves.add(node);
                }
            }

            int minDiff = Integer.MAX_VALUE;

            while (!leaves.isEmpty()) {
                // get a leaf node
                // Node leaf = pq.poll();
                Node leaf = leaves.get(0);

                leaves.remove(0);
                if (leaf.adj.size() <= 0)// last node
                    break;

                int diff = Math.abs((total - leaf.value) - leaf.value);
                if (diff < minDiff)
                    minDiff = diff;

                // combind leaf to it's connection
                Node conn = null;
                for (Node node : leaf.adj.values()) {
                    conn = node;
                }

                conn.value += leaf.value;
                conn.adj.remove(leaf.index);
                if (conn.adj.size() == 1)
                    leaves.add(conn);
                nodes.remove(leaf.index);
            }

            System.out.println(minDiff);
        } catch (FileNotFoundException e) {
            e.printStackTrace();
        } catch (IOException e) {
            // TODO Auto-generated catch block
            e.printStackTrace();
        }
    }
}

Thanks.


Solution

  • The standard Java PriorityQueue does not support update. If you need to remove items, then you have to implement your own minHeap. Call heapify when you removed some items in the heap.

    The implementation and explanation can be found here!