AtCoder Beginner Contest 335 E

https://atcoder.jp/contests/abc335/tasks/abc335_e

繋がっている同じ値のノードは結合してグラフにします。そうすればあとはDPになります。

// Non-Decreasing Colorful Path
#![allow(non_snake_case)]

use std::cmp::max;
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()
}


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

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

fn read_edge() -> Edge {
    let v = read_vec::<Node>();
    (v[0]-1, v[1]-1)
}

struct Graph {
    g: Vec<Vec<Node>>
}

impl Graph {
    fn create(N: usize, edges: &Vec<Edge>) -> Graph {
        let mut g: Vec<Vec<Node>> = vec![vec![]; N];
        for &(u, v) in edges.iter() {
            g[u].push(v)
        }
        Graph { g }
    }
}


struct UnionFind {
    parents: Vec<Node>,
    heights: Vec<i32>,
}

impl UnionFind {
    fn new(N: usize) -> UnionFind {
        let parents: Vec<Node> = (0..N).collect();
        let heights: Vec<i32> = vec![1; N];
        UnionFind { parents, heights }
    }
    
    fn is_root(&self, v: Node) -> bool {
        self.parents[v] == v
    }
    
    fn join(&mut self, u: Node, v: Node) {
        let r1 = self.root(u);
        let r2 = self.root(v);
        if r1 == r2 {
            return  // ここには来ないはず
        }
        
        let h1 = self.heights[r1];
        let h2 = self.heights[r2];
        if h1 <= h2 {   // r2にr1がぶら下がる
            self.parents[r1] = r2;
            self.heights[r2] = max(self.heights[r2], self.heights[r1]+1);
        }
        else {
            self.parents[r2] = r1;
            self.heights[r1] = max(self.heights[r1], self.heights[r2]+1);
        }
    }
    
    fn root(&self, v0: Node) -> Node {
        if self.is_root(v0) {
            v0
        }
        else {
            let p = self.parents[v0];
            self.root(p)
        }
    }
}


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

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

fn collect_same_connected_nodes(A: &Vec<i32>, edges: &Vec<Edge>) -> Vec<Node> {
    let N = A.len();
    let mut uf = UnionFind::new(N);
    for &(u, v) in edges.iter() {
        if A[u] == A[v] {
            uf.join(u, v)
        }
    }
    
    let mut dic: HashMap<Node, Vec<Node>> = HashMap::new();
    for v in 0..N {
        let r = uf.root(v);
        let e = dic.entry(r).or_insert(vec![]);
        (*e).push(v)
    }
    
    let mut roots: Vec<Node> = vec![0; N];
    for (_, vs) in dic.into_iter() {
        let r: Node = if vs.contains(&0) { 0 }
                        else if vs.contains(&(N-1)) { N-1 }
                        else { vs[0] };
        for v1 in vs.into_iter() {
            roots[v1] = r
        }
    }
    roots
}

fn merge_edges(A: &Vec<i32>, edges: &Vec<Edge>,
                                roots: &Vec<Node>) -> Vec<Edge> {
    let mut merged_edges: Vec<Edge> = vec![];
    for &(u, v) in edges.iter() {
        if A[u] < A[v] {
            merged_edges.push((roots[u], roots[v]))
        }
        else if A[u] > A[v] {
            merged_edges.push((roots[v], roots[u]))
        }
    }
    merged_edges
}

fn F(A: Vec<i32>, edges: Vec<Edge>) -> usize {
    let N = A.len();
    if A[0] > A[N-1] {
        return 0
    }
    else if A[0] == A[N-1] {
        return 1
    }
    
    let roots = collect_same_connected_nodes(&A, &edges);
    let merged_edges = merge_edges(&A, &edges, &roots);
    let graph = Graph::create(N, &merged_edges);
    
    let mut B: Vec<(Node, i32)> = (0..N).map(|v| (v, A[v])).collect();
    B.sort_by_key(|&(_, value)| value);
    let mut dp: Vec<usize> = vec![0; N];
    dp[0] = 1;
    for (v, _) in B.into_iter() {
        if dp[v] == 0 {
            continue
        }
        
        for &w in graph.g.get(v).unwrap() {
            dp[w] = max(dp[v] + 1, dp[w])
        }
    }
    dp[N-1]
}

fn main() {
    let (A, edges) = read_input();
    println!("{}", F(A, edges))
}