gistfile1.java

package com.company;

import java.util.ArrayDeque;
import java.util.PriorityQueue;
import java.util.Queue;

public class Main {

    // Start by defining Nodes and Edges (Nodes in a tree have Edges of a certain length to a child Node)
    static class Edge {
        public Edge(int distance, Node child) {
            this.distance = distance;
            this.child = child;
        }
        public int distance;
        public Node child;
    }
    static class Node {
        public Node(String value, Edge[] edges) {
            this.value = value;
            this.edges = edges;
        }
        public Node(String value) {
            this.value = value;
            this.edges = new Edge[0];
        }
        public String value;
        public Edge[] edges;
    }

    // This is just a container for our "return values", that holds each Node along with "total distance to root".
    // It's what goes in our heap, so it needs a "Comparable" implementation that sorts by total distance
    static class NodeWithTotalDistance implements Comparable<NodeWithTotalDistance> {
        public NodeWithTotalDistance(Node node, int distFromRoot) {
            this.node = node;
            this.distFromRoot = distFromRoot;
        }
        public Node node;
        public int distFromRoot;

        public int compareTo(NodeWithTotalDistance other) {
            return distFromRoot - other.distFromRoot;
        }
    }

    // Here's our traversal-by-order-of-total-distance-from-root function.
    // As I suggested in the interview, it follows almost identically the
    // algorithm for depth-first-search using a Queue, except this uses a Heap.
    // This allows you to traverseByTotalDistance by least-total-distance instead of just breadth-first.
    // Note Heap ~ PriorityQueue in Java; also note a heap *is* topologically sorted (the parent always
    // appears before the children in the array) so maybe that's where you were going with that idea?
    // See subsequent function for standard breadth-first algo for comparison.
    static void traverseByTotalDistance(Node root) {
        PriorityQueue<NodeWithTotalDistance> heap = new PriorityQueue<NodeWithTotalDistance>(); // use a heap instead of a queue
        heap.add(new NodeWithTotalDistance(root, 0)); // store total distance along with the node
        while(!heap.isEmpty()) {
            NodeWithTotalDistance nodeDist = heap.remove(); // get next element from heap. it's the least distance one remaining.
            System.out.println(nodeDist.distFromRoot + " " + nodeDist.node.value); // do whatever with it: add to some output list, trigger event, run a callback, yield an iterator, etc.  Here we just write to console.
            for (Edge edge : nodeDist.node.edges) { // add each of children into the heap, calculating their total distance
                heap.add(new NodeWithTotalDistance(edge.child, nodeDist.distFromRoot + edge.distance));
            }
        }
    }
    // Also note that the above is a variant of Dijkstra's algorithm
    // https://en.wikipedia.org/wiki/Dijkstra%27s_algorithm#Practical_optimizations_and_infinite_graphs
    // assuming the goal is unreachable and the graph is acyclic (i.e. a tree), so the two "if" statements in
    // the final "for each" there will always reduce to "true".

    // Here is the breadth-first-search by comparison
    static void traverseBreadthFirst(Node root) {
        Queue<Node> q = new ArrayDeque<Node>();
        q.add(root);
        while(!q.isEmpty()) {
            Node node = q.remove();
            System.out.println(node.value);
            for (Edge edge : node.edges) {
                q.add(edge.child);
            }
        }
    }

    // A test case
    public static void main(String[] args) {
        Node tree = new Node("A", new Edge[] {
                new Edge(1, new Node("B", new Edge[] {
                        new Edge(3, new Node("D")),
                        new Edge(1, new Node("F")) })),
                new Edge(3, new Node("C")) });
        traverseByTotalDistance(tree);
    }
}