Reachability in Functional Graph の解説

なんとなく、sccして、トポロジカル順に並んだ後ろから計算していけば良いことがわかる。

ループと、ループに入れるところを区別するために、ループの頂点は1頂点として持たす。そのためにunionfindを使ってleaderを使用する。

ループ内は要素数^2の到達点がある。

ここまでできたら、sccを用いて、後ろから到達できる頂点を増やしつつbfsしてあげればよい。

fn main() {
    input! {
        n: usize,
        a: [Usize1; n]
    }
    let mut scc = SccGraph::new(n);
    for i in 0..n {
        scc.add_edge(i, a[i]);
    }
    let mut uf = Dsu::new(n);
    let mut scc = scc.scc();
    let mut base = 0;
    for row in scc.iter() {
        for i in 0..row.len() - 1 {
            let now = row[i];
            let next = row[i + 1];
            uf.merge(now, next);
        }
        base += row.len() * row.len();
    }
    let mut map = vec![0; n];
    for row in scc.iter() {
        let leader = uf.leader(row[0]);
        map[leader] = row.len();
    }
    let mut to = vec![HashSet::new(); n];
    for i in 0..n {
        let now = uf.leader(i);
        let target = uf.leader(a[i]);
        if to[target].contains(&now) {
            continue;
        }
        if target == now {
            continue;
        }
        to[target].insert(now);
    }
    let mut vis = vec![false; n];
    let mut ans = vec![0; n];
    scc.reverse();
    for i in 0..scc.len() {
        let i = scc[i][0];
        if vis[i] {
            continue;
        }
        let mut que = VecDeque::new();
        que.push_back(i);
        vis[i] = true;
        while !que.is_empty() {
            let v = que.pop_front().unwrap();
            let count = map[v];
            for &v2 in to[v].iter() {
                ans[v2] += count + ans[v];
                if vis[v2] {
                    continue;
                }
                vis[v2] = true;
                que.push_back(v2);
            }
        }
    }
    let ans = ans.iter().sum::<usize>() + base;
    println!("{}", ans);
}