Sadonkorjuu

use std::io;

fn dijkstra(nodes: &Vec<i32>, edges: &Vec<(usize, usize, i32)>, source: usize, target: usize) -> (i32, i32) {
    let mut dist = vec![std::i32::MAX; nodes.len()];
    dist[source] = 0;

    let mut pq = std::collections::BinaryHeap::new();
    pq.push(std::cmp::Reverse((0, source)));

    while let Some(std::cmp::Reverse((d, u))) = pq.pop() {
        if d > dist[u] {
            continue;
        }

        if u == target {
            return (d, -1);
        } else if nodes[u] == -1 {
            return (d, u as i32);
        }

        for (x, v, w) in edges.iter().filter(|(v, x, _)| (*v == u || *x == u)).map(|(x, v, w)| (*x, *v, *w)) {
            let v = if u == x { v } else { x };
            let alt = d + w;
            if alt < dist[v] {
                dist[v] = alt;
                pq.push(std::cmp::Reverse((alt, v)));
            }
        }
    }

    (-1, -1)
}

fn main() {
    let mut input = String::new();
    io::stdin().read_line(&mut input).unwrap();
    let n = input.trim().parse::<usize>().unwrap();
    let mut nodes = vec![];
    input.clear();
    io::stdin().read_line(&mut input).unwrap();
    let iter = input.trim().split_whitespace();
    for (idx, i) in iter.enumerate() {
        if i == "1" {
            nodes.push(idx as i32);
        } else {
            nodes.push(-1);
        }
    }
    let mut edges = vec![];
    for _ in 1..n {
        input.clear();
        io::stdin().read_line(&mut input).unwrap();
        let mut iter = input.trim().split_whitespace();

        edges.push((iter.next().unwrap().parse::<usize>().unwrap()-1, iter.next().unwrap().parse::<usize>().unwrap()-1, iter.next().unwrap().parse::<i32>().unwrap()));
    }
    let mut sum = 0;
    let mut seen = vec![false; nodes.len()];
    for i in 0..nodes.len() {
        if nodes[i] == -1 {
            continue;
        }
        let mut smallest = std::i32::MAX;
        for j in 0..nodes.len() {
            if nodes[j] != -1 || seen[j] {
                continue;
            }
            let (d, a) = dijkstra(&nodes, &edges, i, j);
            if d < smallest {
                smallest = d;
            }
            if a != -1 {
                seen[a as usize] = true;
            }
        }
        sum += smallest;
        seen = vec![false; nodes.len()];
    }
    println!("{}", sum);
}

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: ACCEPTED

Test 5

Group: 1, 2

Verdict: ACCEPTED

Test 6

Group: 1, 2

Verdict: ACCEPTED

Test 7

Group: 2

Verdict: TIME LIMIT EXCEEDED

Test 8

Group: 1, 2

Verdict: ACCEPTED

Test 9

Group: 2

Verdict: TIME LIMIT EXCEEDED

Test 10

Group: 1, 2

Verdict: ACCEPTED

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: ACCEPTED

Test 16

Group: 1, 2

Verdict: ACCEPTED

Test 17

Group: 1, 2

Verdict: ACCEPTED

Test 18

Group: 1, 2

Verdict: TIME LIMIT EXCEEDED

Test 19

Group: 1, 2

Verdict: ACCEPTED

Test 20

Group: 1, 2

Verdict: ACCEPTED

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: ACCEPTED

Test 29

Group: 2

Verdict: TIME LIMIT EXCEEDED

Test 30

Group: 1, 2

Verdict: ACCEPTED

Test 31

Group: 2

Verdict: TIME LIMIT EXCEEDED