AtCoder Beginner Contest 438 F

https://atcoder.jp/contests/abc438/tasks/abc438_f

数える順番を変える、じゃないですが、数え方を変えます。すなわち、問題文はパスを列挙してそのパスに含まれない最小値を足しますが、最小値を固定してそうなるパスを数えます。例えば入力例2で最小値0なら、0を含まないパスを数えればいいから、0をrootとする木を考えると、ノードが1個ぶら下がっている枝と3個ぶら下がっている枝があるので、それぞれの枝の中だけで端点を取ればいいので、1+6=7個になります。これなら数えられそうです。
もう少し詳細に考えましょう。最小が1になる場合を考えます。0は通っているけど1は通っていないパスを数えることになりますが、それより0を通っているパスの数と0と1を通っているパスの数を数えて引き算するほうが考えやすいでしょう。このように0から順に加えていったノードを全て通るパスを数えるとよいです。そういうパスが一つもなくなったらそこで終わりです。
0と1のあと2を追加するなら、2が1より下にあるか、2が0と1の間にあるか、2が1と違う枝にあるかです。しかし、2が1より下にあるかすぐにはわからないですよね。ノードIDがバラバラなので、子の範囲の情報がないのでダメです。
それならノードIDを並び替えればよいです。rootは0として、子は親よりIDが大きいとして木のノードを振りなおします。そして、元のノードIDから振りなおしたIDへのテーブルを用意しておけばよいです。
最後に、木が1列になっていれば、ノード数を足します。

// Sum of Mex
#![allow(non_snake_case)]


//////////////////// 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: Vec<Node> = read_vec();
    (v[0], v[1])
}

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

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


//////////////////// Vertex ////////////////////

struct Vertex {
    first: usize,
    last: usize,
    children: Vec<usize>
}

impl Vertex {
    fn is_leaf(&self) -> bool {
        self.children.is_empty()
    }
    
    fn contains(&self, ID: usize) -> bool {
        self.first <= ID && ID < self.last
    }
    
    fn count(&self) -> usize {
        self.last - self.first
    }
    
    fn create(ID: usize) -> Vertex {
        Vertex { first: ID, last: 0, children: vec![] }
    }
}


//////////////////// Tree ////////////////////

struct Tree {
    // 並び変えたIDはusize
    v: Vec<Vertex>,
    a: Vec<usize>,      // 元のID -> 並び変えたID
}

impl Tree {
    // 0を含むパスの数
    fn count_paths_with_0(&self) -> usize {
        // 0からのブランチが含むnodeの個数をn1, n2, ...とすると
        // n1n2 + n1n3 + ... = ((n1+n2+...)^2 - (n1^2+n2^2+...))/2
        let N = self.v.len();
        let mut s = (N - 1) * (N - 1);
        for &child in self.v[0].children.iter() {
            let n = self.v[child].count();
            s = s + n * 2 - n * n
        }
        s / 2 + 1   // (0, 0)も加える
    }
    
    // IDをそれ以下に含むのは何番目の子か
    fn find_branch(&self, ID: usize) -> usize {
        let children = &self.v[0].children;
        let mut first: usize = 0;
        let mut last = children.len();
        while last - first > 1 {
            let mid = (first + last) / 2;
            if self.v[children[mid]].first > ID {
                last = mid
            }
            else {
                first = mid
            }
        }
        first
    }
    
    // IDを含む0のブランチ以外でnodeをいくつ持つか
    fn count_without_branch(&self, ID: usize) -> usize {
        // k番目のbranchにIDが含まれる
        let k = self.find_branch(ID);
        let child_ID = self.v[0].children[k];
        // そのブランチ以外のノードの個数
        self.v.len() - self.v[child_ID].count()
    }
    
    // 0を含む2つのIDを通るパスを数える
    fn count_paths_with_0_and_other_node(&self, ID: usize) -> usize {
        // IDを含む0のブランチ以外でnodeをいくつ持つか
        let n1 = self.count_without_branch(ID);
        // ID以下のノードの個数
        let n2 = self.v[ID].count();
        n1 * n2
    }
    
    // 0をはさむ2つのIDを通るパスを数える
    fn count_paths_with_three_nodes(&self, ID1: usize, ID2: usize) -> usize {
        let n1 = self.v[ID1].count();
        let n2 = self.v[ID2].count();
        n1 * n2
    }
    
    fn is_linear(&self) -> bool {
        if self.v[0].children.len() > 2 {
            return false
        }
        
        for &child in self.v[0].children.iter() {
            let mut v = child;
            while !self.v[v].is_leaf() {
                if self.v[v].children.len() != 1 {
                    return false
                }
                v = self.v[v].children[0]
            }
        }
        true
    }
    
    fn count(&self) -> usize {
        let mut counter: usize = 0;
        let N = self.v.len();
        let mut IDs: Vec<usize> = vec![];
        let mut s0 = T(N);  // IDsを全て通るパスの数
        for v in 0..N {
            let ID = self.a[v];
            // IDsにIDを加えたnodeを全て通るパスを数える
            let s: usize;
            if IDs.len() == 0 {
                s = self.count_paths_with_0();
                IDs.push(ID)
            }
            else if IDs.len() == 1 {
                // 0のみ
                s = self.count_paths_with_0_and_other_node(ID);
                IDs.push(ID)
            }
            else if IDs.len() == 2 {
                // 0と別のID1のみ
                let ID1 = IDs[1];
                let k = self.find_branch(ID);
                let child_node = &self.v[self.v[0].children[k]];
                if self.v[ID].contains(ID1) {
                    // IDは0とID1の間にあるので、パスの数は変わらず
                    s = s0
                }
                else if self.v[ID1].contains(ID) {
                    // IDはID1の下
                    s = self.count_paths_with_0_and_other_node(ID);
                    IDs[1] = ID
                }
                else if !child_node.contains(ID1) {
                    // IDが0から見てID1と違うブランチにある
                    let n1 = self.v[ID1].count();
                    let n2 = self.v[ID].count();
                    s = n1 * n2;
                    IDs.push(ID)
                }
                else {
                    // IDが0のID1を含むブランチにあって、ID1と違うパスにあるとき
                    // 全て通るパスはないので、vより大きい最小値になることはない
                    counter += v * s0;
                    break
                }
            }
            else {
                // すでに3つのノードを通るのが必須
                let ID1 = IDs[1];
                let ID2 = IDs[2];
                if self.v[ID].contains(ID1) {
                    // IDは0とID1の間
                    s = s0
                }
                else if self.v[ID1].contains(ID) {
                    // IDはID1の下
                    s = self.count_paths_with_three_nodes(ID, ID2);
                    IDs[1] = ID
                }
                else if self.v[ID].contains(ID2) {
                    // IDは0とID2の間
                    s = s0
                }
                else if self.v[ID2].contains(ID) {
                    // IDはID2の下
                    s = self.count_paths_with_three_nodes(ID1, ID);
                    IDs[2] = ID
                }
                else {
                    counter += v * s0;
                    break
                }
            }
            counter += v * (s0 - s);
            s0 = s
        }
        
        if self.is_linear() {
            counter += N
        }
        counter
    }
    
    // Vertexを初期化して、その下も作る
    // 次のIDを返す
    fn initialize_subtree(&mut self, prev_v: Node, cur_v: Node,
                                ID0: usize, graph: &Graph) -> usize {
        self.a[cur_v] = ID0;
        let mut ID = ID0 + 1;
        for &next_v in graph.g[cur_v].iter() {
            if next_v != prev_v {
                let prev_ID = ID;
                ID = self.initialize_subtree(cur_v, next_v, ID, graph);
                self.v[ID0].children.push(prev_ID)
            }
        }
        self.v[ID0].last = ID;
        ID
    }
    
    fn create(graph: Graph) -> Tree {
        let N = graph.len();
        let v: Vec<Vertex> = (0..N).map(|ID| Vertex::create(ID)).collect();
        let a: Vec<Node> = vec![0; N];
        let mut tree = Tree { v, a };
        tree.initialize_subtree(0, 0, 0, &graph);
        tree
    }
}


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

fn read_input() -> Vec<Edge> {
    let N: usize = read();
    let edges: Vec<Edge> = (0..N-1).map(|_| read_edge()).collect();
    edges
}

fn T(n: usize) -> usize {
    n * (n + 1) / 2
}

fn F(edges: Vec<Edge>) -> usize {
    let N = edges.len() + 1;
    let graph = Graph::create(N, edges);
    let tree = Tree::create(graph);
    tree.count()
}

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