Sadonkorjuu

//package Harvest;

import java.util.*;
import java.util.Scanner;

public class Harvest {    
    public static void main(String[] args) {
        try (Scanner sc = new Scanner(System.in)) {
            int noOfCities = Integer.parseInt(sc.nextLine());
            
            String [] cityStatus = sc.nextLine().split(" ");

            ArrayList<Vertex> cities = new ArrayList<>();

            Graph cityNetwork = new Graph(true,false);

            for (int cityIndex = 1; cityIndex<=noOfCities; cityIndex++)
            {
                cities.add(cityIndex-1,cityNetwork.addVertex(String.valueOf(cityIndex)));
            }

            for (int counter = 1; counter<noOfCities; counter++)
            {
                String [] networkInformation = sc.nextLine().split(" ");
                cityNetwork.addEdge(cities.get(Integer.parseInt(networkInformation[0])-1),cities.get(Integer.parseInt(networkInformation[1])-1),Integer.parseInt(networkInformation[2]));
            }

            ArrayList<Integer> shortestDistances = new ArrayList<>();

            int shortestDistance = 0;
            for (Vertex city: cities){
                if (Integer.parseInt(cityStatus[cities.indexOf(city)])==1)
                {
                    for (Vertex otherCity: cities)
                    {
                        if (Integer.parseInt(cityStatus[cities.indexOf(otherCity)])==0)
                        {
                            shortestDistance = Dijkstra.shortestPathBetween(cityNetwork,city, otherCity);
                        }
                    }
                    shortestDistances.add(shortestDistance);
                }
            }

            int finalSum = 0;
            for(int distance:shortestDistances)
            {
                finalSum=finalSum + distance;
            }

            System.out.println(finalSum);

        } catch (NumberFormatException e) {
            e.printStackTrace();
        }
    }

}

class Dijkstra{

    public static Dictionary[] dijkstra(Graph g, Vertex startingVertex){
        Dictionary<String,Integer> distances = new Hashtable<>();
        Dictionary<String,Vertex> previous = new Hashtable<>(); 

        PriorityQueue<QueueObject> queue = new PriorityQueue<QueueObject>();

        queue.add(new QueueObject(startingVertex, 0));

        for (Vertex v: g.getVertices()){
            if (v != startingVertex){
                distances.put(v.getData(), Integer.MAX_VALUE);
            }
            previous.put(v.getData(), new Vertex(null));
        }

        distances.put(startingVertex.getData(), 0);

        while (queue.size()!=0)
        {
            Vertex current = queue.poll().vertex;
            for (Edge e: current.getEdges()){
                Integer alternative = distances.get(current.getData()) + e.getWeight();
                String neighbourValue = e.getEnd().getData();
                if (alternative<distances.get(neighbourValue)){
                    distances.put(neighbourValue, alternative);
                    previous.put(neighbourValue,current);
                    queue.add(new QueueObject(e.getEnd(), distances.get(neighbourValue)));
                }
            }
        }
        return new Dictionary[] {distances, previous};
    }

    public static void dijkstraResultPrinter(Dictionary[] d)
    {
        System.out.println("Distances:\n");
        for (Enumeration keys = d[0].keys(); keys.hasMoreElements();){
            String nextKey = keys.nextElement().toString();
            System.out.println(nextKey + ": " + d[0].get(nextKey));
        }
        /*System.out.println("\nPrevious:\n");
        for (Enumeration keys = d[1].keys(); keys.hasMoreElements();){
            String nextKey = keys.nextElement().toString();
            Vertex nextVertex = (Vertex) d[1].get(nextKey);
            System.out.println(nextKey + ": " + nextVertex.getData());
        }*/
    }

    public static int shortestPathBetween(Graph g, Vertex startingVertex, Vertex targetVertex)
    {
        Dictionary[] dijkstraDictionaries = dijkstra(g, startingVertex);
        Dictionary distances = dijkstraDictionaries[0];

        Integer distance = (Integer) distances.get(targetVertex.getData());
        //System.out.println("Shortest distance between " + startingVertex.getData() + " and " + targetVertex.getData());
        return distance;
    }

}

class Vertex
{
    private String data;
    private ArrayList<Edge> edges;

    public Vertex(String inputData)
    {
        this.data = inputData;
        this.edges = new ArrayList<Edge>();
    }

    public void addEdge(Vertex endVertex, Integer weight)
    {
        this.edges.add(new Edge(this, endVertex, weight));
    }

    public void removeEdge(Vertex endVertex)
    {
        this.edges.removeIf(edge -> edge.getEnd().equals(endVertex));
    }

    public String getData()
    {
        return this.data;
    }

    public ArrayList<Edge> getEdges()
    {
        return this.edges;
    }

    public void print(boolean showWeight)
    {
        String message = "";

        if (this.edges.size() == 0)
        {
            System.out.println(this.data + " --->");
            return;
        }

        for (int i = 0; i < this.edges.size(); i++)
        {
            if (i == 0)
                message += this.edges.get(i).getStart().data + " ---> ";

            message += this.edges.get(i).getEnd().data;

            if (showWeight)
                message += " (" + this.edges.get(i).getWeight() + ")";

            if (i != this.edges.size() - 1)
                message += ", ";

            System.out.println(message);
        }
    }
}

class QueueObject implements Comparable<QueueObject>{
    public Vertex vertex;
    public int priority;

    public QueueObject(Vertex v, int p)
    {
        this.vertex=v;
        this.priority = p;
    }

    @Override
    public int compareTo(QueueObject o)
    {
        if (this.priority == o.priority){
            return 0;
        }
        else if (this.priority>o.priority)
        {
            return 1;
        }
        else{
            return -1;
        }
    }
}

class Queue{
    
    public LinkedList queue;
    public int size;
    static final int DEFAULT_MAX_SIZE = Integer.MAX_VALUE;
    public int maxSize;

    public Queue()
    {
        this(DEFAULT_MAX_SIZE);
    }

    public Queue(int maxSize) {
        this.queue = new LinkedList();
        this.size = 0;
        this.maxSize = maxSize;
    }
    
    public boolean hasSpace ()
    {
        return this.size<this.maxSize;
    }

    public boolean isEmpty()
    {
        return this.size == 0;
    }

    public void enqueue(Vertex data)
    {
        if (this.hasSpace())
        {
            this.queue.addToHead(data);
            this.size++;
        } else {
            throw new Error ("Queue is full!");
        }
    }

    public Vertex dequeue()
    {
        if (!this.isEmpty()){
            Vertex data = this.queue.removeHead();
            this.size--;
            return data;
        } else {
            throw new Error ("Queue is empty!");
        }
    }

    public Vertex peak()
    {
        if (this.isEmpty())
            return null;
        else
            return this.queue.head.data;    }

}

class Node {
    public Vertex data;
    private Node next;

    public Node (Vertex data)
    {
        this.data = data;
        this.next=null;
    }
    public Node getNextNode() {
        return this.next;
    }

    public void setNextNode(Node next) {
        this.next = next;
    }
}

class LinkedList {
    public Node head;

    public LinkedList(){
        this.head = null;
    }

    public void addToHead(Vertex data)
    {
        Node newNode = new Node (data);
        Node currentNode = this.head;

        this.head = newNode;
        if (currentNode!=null){
            this.head.setNextNode(currentNode);
        }
    }

    public void addToTail(Vertex data)
    {
        Node tail = this.head;
        if (tail == null)
            this.head = new Node(data);
        else{
            while(tail.getNextNode() != null)
                tail = tail.getNextNode();
            tail.setNextNode(new Node(data));
        }
    }

    public Vertex removeHead(){
        Node removedHead = this.head;
        if (removedHead ==null ){
            return null;
        }
        this.head = removedHead.getNextNode();
        return removedHead.data;    
    }

    public String print()
    {
        String output = "<head> ";
        Node currenNode = this.head;

        while (currenNode != null)
        {
            output+=currenNode.data + " ";
            currenNode = currenNode.getNextNode();
        }

        output += "<tail>";
        System.out.println(output);
        return output;
    }
}

class Graph
{
    private ArrayList<Vertex> vertices;
    private boolean isWeighted;
    private boolean isDirected;

    public Graph(boolean inputIsWeighted, boolean inputIsDirected)
    {
        this.vertices = new ArrayList<Vertex>();
        this.isWeighted = inputIsWeighted;
        this.isDirected = inputIsDirected;
    }

    public Vertex addVertex (String data)
    {
        Vertex newVertex = new Vertex(data);
        this.vertices.add(newVertex);
        return newVertex;
    }

    public void addEdge (Vertex vertex1, Vertex vertex2, Integer weight)
    {
        if(!this.isWeighted)
            weight=null;

        vertex1.addEdge(vertex2, weight);
        if (!this.isDirected)
            vertex2.addEdge(vertex1, weight);
    }

    public void removeEdge(Vertex vertex1, Vertex vertex2)
    {
        vertex1.removeEdge(vertex2);
        if (!this.isDirected)
            vertex2.removeEdge(vertex1);
    }

    public void removeVertex(Vertex vertex)
    {
        this.vertices.remove(vertex);
    }

    public ArrayList<Vertex> getVertices()
    {
        return this.vertices;
    }

    public boolean isWeighted()
    {
        return this.isWeighted;
    }

    public boolean isDirected()
    {
        return this.isDirected;
    }

    public Vertex getVertexByValue(String value)
    {
        for (Vertex v: this.vertices)
            if (Objects.equals(v.getData(), value))
                return v;
        return null;
    }

    public void print()
    {
        for (Vertex v: this.vertices)
            v.print(isWeighted);
    }

    public Vertex getStartingVertex()
    {
        return this.getVertices().get(0);
    }
}

class Edge
{
    private Vertex start;
    private Vertex end;
    private Integer weight;

    public Edge (Vertex startV, Vertex endV, Integer inputWeight)
    {
        this.start = startV;
        this.end = endV;
        this.weight = inputWeight;
    }

    public Vertex getStart()
    {
        return this.start;
    }

    public Vertex getEnd()
    {
        return this.end;
    }

    public Integer getWeight()
    {
        return this.weight;
    }
}

Test 1

Group: 1, 2

Verdict: ACCEPTED

Test 2

Group: 1, 2

Verdict: ACCEPTED

Test 3

Group: 1, 2

Verdict: ACCEPTED

Test 4

Group: 1, 2

Verdict: WRONG ANSWER

Test 5

Group: 1, 2

Verdict: WRONG ANSWER

Test 6

Group: 1, 2

Verdict: TIME LIMIT EXCEEDED

Test 7

Group: 2

Verdict: TIME LIMIT EXCEEDED

Test 8

Group: 1, 2

Verdict: TIME LIMIT EXCEEDED

Test 9

Group: 2

Verdict: TIME LIMIT EXCEEDED

Test 10

Group: 1, 2

Verdict: TIME LIMIT EXCEEDED

Test 11

Group: 2

Verdict: TIME LIMIT EXCEEDED

Test 12

Group: 2

Verdict: TIME LIMIT EXCEEDED

Test 13

Group: 2

Verdict: TIME LIMIT EXCEEDED

Test 14

Group: 2

Verdict: TIME LIMIT EXCEEDED

Test 15

Group: 1, 2

Verdict: TIME LIMIT EXCEEDED

Test 16

Group: 1, 2

Verdict: TIME LIMIT EXCEEDED

Test 17

Group: 1, 2

Verdict: TIME LIMIT EXCEEDED

Test 18

Group: 1, 2

Verdict: TIME LIMIT EXCEEDED

Test 19

Group: 1, 2

Verdict: TIME LIMIT EXCEEDED

Test 20

Group: 1, 2

Verdict: TIME LIMIT EXCEEDED

Test 21

Group: 2

Verdict: TIME LIMIT EXCEEDED

Test 22

Group: 2

Verdict: TIME LIMIT EXCEEDED

Test 23

Group: 2

Verdict: TIME LIMIT EXCEEDED

Test 24

Group: 1, 2

Verdict: TIME LIMIT EXCEEDED

Test 25

Group: 2

Verdict: TIME LIMIT EXCEEDED

Test 26

Group: 1, 2

Verdict: TIME LIMIT EXCEEDED

Test 27

Group: 2

Verdict: TIME LIMIT EXCEEDED

Test 28

Group: 1, 2

Verdict: TIME LIMIT EXCEEDED

Test 29

Group: 2

Verdict: TIME LIMIT EXCEEDED

Test 30

Group: 1, 2

Verdict: TIME LIMIT EXCEEDED

Test 31

Group: 2

Verdict: TIME LIMIT EXCEEDED