UVa 908 - Re-connecting Computer Sites

// UVa 908 - Re-connecting Computer Sites

import java.util.*;

public class Main {

    public static void main(String[] args) {

        Scanner scanner = new Scanner(System.in);

        int testCase = 0;
        while (scanner.hasNextInt()) {
            testCase++;
            if (testCase > 1)
                System.out.println();

            Mst mst = new Mst();

            int n = scanner.nextInt();
            for (int i = 1; i < n; i++) {
                int a = scanner.nextInt(), b = scanner.nextInt(), c = scanner.nextInt();
                mst.addEdge(a, b, c);
            }
            System.out.println(mst.getCurrentCost());

            int k = scanner.nextInt();
            for (int i = 0; i < k; i++) {
                int a = scanner.nextInt(), b = scanner.nextInt(), c = scanner.nextInt();
                mst.offerNewEdge(a, b, c);
            }
            System.out.println(mst.getCurrentCost());

            int m = scanner.nextInt();
            for (int i = 0; i < m; i++) {
                scanner.nextInt();
                scanner.nextInt();
                scanner.nextInt();
            }
        }
    }

    private static class Mst {

        private int currentCost = 0;
        private Map<Integer, Map<Integer, Edge>> matrix = new HashMap<>();

        public int getCurrentCost() {
            return currentCost;
        }

        public int offerNewEdge(int from, int to, int cost) {
            Edge costlyEdge = getMostCostlyEdgeInPath(from, to);
            if (costlyEdge.cost > cost) {
                removeEdge(costlyEdge);
                addEdge(from, to, cost);
            }
            return currentCost;
        }


        private Edge getMostCostlyEdgeInPath(int from, int to) {
            Queue<EdgePair> q = new LinkedList<>();
            q.add(new EdgePair(
                    new Edge(-1, from, 0),
                    new Edge(-1, from, 0)
            ));
            while (!q.isEmpty()) {
                EdgePair head = q.poll();
                int node = head.last.to;
                for (Integer next : matrix.get(node).keySet()) {
                    if (next == head.last.from) continue;
                    Edge edge = matrix.get(node).get(next);
                    Edge newCostlyEdge = edge.cost > head.costly.cost
                            ? edge : head.costly;
                    if (next == to)
                        return newCostlyEdge;
                    q.add(new EdgePair(
                            edge,
                            newCostlyEdge
                    ));
                }
            }
            return null;
        }

        private static class EdgePair {
            public Edge last, costly;

            public EdgePair(Edge last, Edge costly) {
                this.last = last;
                this.costly = costly;
            }
        }

        private static class Edge {

            public final int from, to, cost;

            public Edge(int from, int to, int cost) {
                this.from = from;
                this.to = to;
                this.cost = cost;
            }
        }

        private void removeEdge(Edge edge) {
            matrix.get(edge.from).remove(edge.to);
            matrix.get(edge.to).remove(edge.from);
            currentCost -= edge.cost;
        }

        public void addEdge(int from, int to, int cost) {
            addToMatrix(from, to, cost);
            addToMatrix(to, from, cost);
            currentCost += cost;
        }

        private void addToMatrix(int a, int b, int c) {
            if (!matrix.containsKey(a))
                matrix.put(a, new HashMap<>());
            matrix.get(a).put(b, new Edge(a, b, c));
        }

    }


}