AtCoder Beginner Contest 369 E

https://atcoder.jp/contests/abc369/tasks/abc369_e

ワーシャルフロイド法で各ノード間の最短距離を求めておいて、クエリーの端を並び替えてその中でDPで橋の両端を入れ替えて最短距離を求めます。

// Bonus EXP
#![allow(non_snake_case)]

use std::collections::HashMap;


//////////////////// library ////////////////////

fn read<T: std::str::FromStr>() -> T {
    let mut line = String::new();
    std::io::stdin().read_line(&mut line).ok();
    line.trim().parse().ok().unwrap()
}

fn read_vec<T: std::str::FromStr>() -> Vec<T> {
    read::<String>().split_whitespace()
            .map(|e| e.parse().ok().unwrap()).collect()
}

fn permutations(a: Vec<Edge>) -> Vec<Vec<Edge>> {
    let N = a.len();
    if N == 1 {
        vec![a.to_vec()]
    }
    else {
        let mut bs: Vec<Vec<Edge>> = vec![];
        for i in 0..N {
            let v: Vec<Edge> = (0..N).filter(|&j| j != i).
                                            map(|j| a[j]).collect();
            let bs_sub = permutations(v);
            for mut b in bs_sub.into_iter() {
                b.insert(0, a[i]);
                bs.push(b)
            }
        }
        bs
    }
}


//////////////////// Graph ////////////////////

type Node = usize;
type Weight = i64;
type Edge = (Node, Node, Weight);

fn read_edge() -> Edge {
    let v: Vec<usize> = read_vec();
    (v[0] - 1, v[1] - 1, v[2] as Weight)
}

struct Graph {
    g: Vec<Vec<(Node, Weight)>>
}

impl Graph {
    fn len(&self) -> usize {
        self.g.len()
    }
    
    fn reduce(&self) -> Graph {
        let N = self.len();
        let mut g: Vec<Vec<(Node, Weight)>> = vec![vec![]; N];
        for (v, e) in self.g.iter().enumerate() {
            let mut m: HashMap<Node, Weight> = HashMap::new();
            for (u, t) in e.into_iter() {
                match m.get(u) {
                    Some(t1) => if t1 > t { m.insert(*u, *t); },
                    None     => { m.insert(*u, *t); },
                }
            }
            for (u, t) in m.into_iter() {
                g[v].push((u, t))
            }
        }
        Graph { g }
    }
    
    fn FloydWarshall(&self) -> Vec<Vec<Weight>> {
        let N = self.len();
        let mut W = vec![vec![-1; N]; N];
        for (i, v) in self.g.iter().enumerate() {
            W[i][i] = 0;
            for &(j, t) in v.iter() {
                W[i][j] = t
            }
        }
        
        for k in 0..N {
            for i in 0..N {
                for j in 0..N {
                    if W[i][k] == -1 || W[k][j] == -1 {
                        continue
                    }
                    let w = W[i][k] + W[k][j];
                    if W[i][j] == -1 || w < W[i][j] {
                        W[i][j] = w
                    }
                }
            }
        }
        W
    }
    
    fn create(N: usize, edges: &Vec<Edge>) -> Graph {
        let mut g: Vec<Vec<(Node, Weight)>> = vec![vec![]; N];
        for &(U, V, T) in edges.iter() {
            g[U].push((V, T));
            g[V].push((U, T))
        }
        Graph { g }
    }
}


//////////////////// Query ////////////////////

type Query = Vec<usize>;

fn read_query() -> Query {
    let _K: usize = read();
    let B: Query = read_vec::<usize>().into_iter().map(|v| v-1).collect();
    B
}


//////////////////// process ////////////////////

fn read_input() -> (usize, Vec<Edge>, usize) {
    let v: Vec<usize> = read_vec();
    let (N, M) = (v[0], v[1]);
    let edges: Vec<Edge> = (0..M).map(|_| read_edge()).collect();
    let Q: usize = read();
    (N, edges, Q)
}

use std::cmp::min;

type DP = (Weight, Weight);

fn initialize_dp(edge: Edge, W: &Vec<Vec<Weight>>) -> DP {
    let w1 = W[0][edge.1] + edge.2;
    let w2 = W[0][edge.0] + edge.2;
    (w1, w2)
}

fn update_dp(edge1: Edge, edge2: Edge, W: &Vec<Vec<Weight>>, dp: DP) -> DP {
    let w11 = dp.0 + W[edge1.0][edge2.1] + edge2.2;
    let w12 = dp.0 + W[edge1.0][edge2.0] + edge2.2;
    let w21 = dp.1 + W[edge1.1][edge2.1] + edge2.2;
    let w22 = dp.1 + W[edge1.1][edge2.0] + edge2.2;
    (min(w11, w21), min(w12, w22))
}

fn calc_dist(edges: Vec<Edge>, W: &Vec<Vec<Weight>>) -> Weight {
    let N = W.len();
    let mut dp = initialize_dp(edges[0], W);
    for i in 1..edges.len() {
        dp = update_dp(edges[i-1], edges[i], W, dp)
    }
    let last_edge = edges[edges.len()-1];
    min(dp.0 + W[last_edge.0][N-1], dp.1 + W[last_edge.1][N-1])
}

fn find_shortest(es: Vec<Edge>, W: &Vec<Vec<Weight>>) -> Weight {
    let perms = permutations(es);
    perms.into_iter().map(|es1| calc_dist(es1, W)).into_iter().min().unwrap()
}

fn F(N: usize, edges: Vec<Edge>, Q: usize) {
    let graph = Graph::create(N, &edges);
    let subg = graph.reduce();
    let W = subg.FloydWarshall();
    for _ in 0..Q {
        let bridges = read_query();
        let es: Vec<Edge> = bridges.into_iter().map(|i| edges[i]).collect();
        println!("{}", find_shortest(es, &W))
    }
}

fn main() {
    let (N, edges, Q) = read_input();
    F(N, edges, Q)
}