#724436 (Rust) No.1787 Do Use Dynamic Tree

提出ソース

結果

問題	No.1787 Do Use Dynamic Tree
コンテスト
ユーザー	akakimidori
提出日時	2021-12-16 19:53:47
言語	Rust (1.83.0 + proconio)
結果	AC
実行時間	1,429 ms / 10,000 ms
コード長	19,936 bytes
コンパイル時間	14,277 ms
コンパイル使用メモリ	388,212 KB
実行使用メモリ	92,744 KB
最終ジャッジ日時	2024-09-13 17:39:45
合計ジャッジ時間	32,444 ms
ジャッジサーバーID （参考情報）	judge2 / judge1

このコードへのチャレンジ
（要ログイン）

ファイルパターン	結果
sample	AC * 2
other	AC * 38

権限があれば一括ダウンロードができます

コンパイルメッセージ

warning: unused variable: `it`
   --> src/main.rs:519:10
    |
519 |     for (it, (ans, (a, b))) in ans.iter_mut().zip(ask).enumerate() {
    |          ^^ help: if this is intentional, prefix it with an underscore: `_it`
    |
    = note: `#[warn(unused_variables)]` on by default

warning: function `rand_memory` is never used
  --> src/main.rs:65:4
   |
65 | fn rand_memory() -> usize {
   |    ^^^^^^^^^^^
   |
   = note: `#[warn(dead_code)]` on by default

warning: function `rand` is never used
  --> src/main.rs:69:4
   |
69 | fn rand() -> usize {
   |    ^^^^

warning: function `shuffle` is never used
  --> src/main.rs:82:4
   |
82 | fn shuffle<T>(a: &mut [T]) {
   |    ^^^^^^^

warning: function `naive` is never used
   --> src/main.rs:600:4
    |
600 | fn naive(e: Vec<(usize, usize)>, ask: Vec<(usize, usize)>) -> Vec<usize> {
    |    ^^^^^

ソースコード

raw source code

//---------- begin union_find ----------
pub struct DSU {
    p: Vec<i32>,
}
impl DSU {
    pub fn new(n: usize) -> DSU {
        assert!(n < std::i32::MAX as usize);
        DSU { p: vec![-1; n] }
    }
    pub fn init(&mut self) {
        self.p.iter_mut().for_each(|p| *p = -1);
    }
    pub fn root(&self, mut x: usize) -> usize {
        assert!(x < self.p.len());
        while self.p[x] >= 0 {
            x = self.p[x] as usize;
        }
        x
    }
    pub fn same(&self, x: usize, y: usize) -> bool {
        assert!(x < self.p.len() && y < self.p.len());
        self.root(x) == self.root(y)
    }
    pub fn unite(&mut self, x: usize, y: usize) -> Option<(usize, usize)> {
        assert!(x < self.p.len() && y < self.p.len());
        let mut x = self.root(x);
        let mut y = self.root(y);
        if x == y {
            return None;
        }
        if self.p[x] > self.p[y] {
            std::mem::swap(&mut x, &mut y);
        }
        self.p[x] += self.p[y];
        self.p[y] = x as i32;
        Some((x, y))
    }
    pub fn parent(&self, x: usize) -> Option<usize> {
        assert!(x < self.p.len());
        let p = self.p[x];
        if p >= 0 {
            Some(p as usize)
        } else {
            None
        }
    }
    pub fn sum<F>(&self, mut x: usize, mut f: F) -> usize
    where
        F: FnMut(usize),
    {
        while let Some(p) = self.parent(x) {
            f(x);
            x = p;
        }
        x
    }
    pub fn size(&self, x: usize) -> usize {
        assert!(x < self.p.len());
        let r = self.root(x);
        (-self.p[r]) as usize
    }
}
//---------- end union_find ----------

fn rand_memory() -> usize {
    Box::into_raw(Box::new("I hope this is a random number")) as usize
}

fn rand() -> usize {
    static mut X: usize = 0;
    unsafe {
        if X == 0 {
            X = rand_memory();
        }
        X ^= X << 13;
        X ^= X >> 17;
        X ^= X << 5;
        X
    }
}

fn shuffle<T>(a: &mut [T]) {
    for i in 1..a.len() {
        let p = rand() % (i + 1);
        a.swap(i, p);
    }
}
// ---------- begin chmin, chmax ----------
pub trait ChangeMinMax {
    fn chmin(&mut self, x: Self) -> bool;
    fn chmax(&mut self, x: Self) -> bool;
}

impl<T: PartialOrd> ChangeMinMax for T {
    fn chmin(&mut self, x: Self) -> bool {
        *self > x && {
            *self = x;
            true
        }
    }
    fn chmax(&mut self, x: Self) -> bool {
        *self < x && {
            *self = x;
            true
        }
    }
}
// ---------- end chmin, chmax ----------
// ---------- begin Heavy-Light decomposition ----------
pub struct HLD {
    size: usize,
    edge: Vec<(usize, usize)>,
    child: Vec<Vec<usize>>,
    path_root: Vec<usize>,
    parent: Vec<usize>,
    left: Vec<usize>,
    right: Vec<usize>,
    inverse: Vec<usize>,
    path_range: Vec<(usize, usize)>,
}

impl HLD {
    pub fn new(size: usize) -> Self {
        assert!(size <= 10usize.pow(8));
        HLD {
            size: size,
            edge: Vec::with_capacity(size - 1),
            child: Vec::new(),
            path_root: Vec::new(),
            parent: Vec::new(),
            left: Vec::new(),
            right: Vec::new(),
            inverse: Vec::new(),
            path_range: vec![],
        }
    }
    pub fn add_edge(&mut self, a: usize, b: usize) {
        assert!(a != b && a < self.size && b < self.size);
        self.edge.push((a, b));
    }
    pub fn build(&mut self, root: usize) {
        assert!(self.edge.len() + 1 == self.size);
        let size = self.size;
        let mut cnt = vec![0; size];
        for &(a, b) in self.edge.iter() {
            cnt[a] += 1;
            cnt[b] += 1;
        }
        let mut child = cnt
            .into_iter()
            .map(|c| Vec::with_capacity(c))
            .collect::<Vec<_>>();
        for &(a, b) in self.edge.iter() {
            child[a].push(b);
            child[b].push(a);
        }
        let mut parent = vec![size; size];
        let mut q = Vec::with_capacity(size);
        q.push(root);
        parent[root] = root;
        for i in 0..size {
            let v = q[i];
            for u in child[v].clone() {
                assert!(parent[u] == size);
                parent[u] = v;
                let k = child[u].iter().position(|e| *e == v).unwrap();
                child[u].remove(k);
                q.push(u);
            }
        }
        let mut sum = vec![1; size];
        for &v in q.iter().rev() {
            let child = &mut child[v];
            if !child.is_empty() {
                let mut max = (0, 0);
                for (i, &u) in child.iter().enumerate() {
                    sum[v] += sum[u];
                    max = std::cmp::max(max, (sum[u], i));
                }
                child.swap(0, max.1);
            }
        }
        let mut path_root = (0..size).collect::<Vec<_>>();
        let mut left = vec![0; size];
        let mut right = vec![0; size];
        let mut path_range = vec![(size, 0); size];
        let mut dfs = vec![(root, false)];
        let mut id = 0;
        while let Some((v, end)) = dfs.pop() {
            if end {
                right[v] = id;
                continue;
            }
            path_range[path_root[v]].0.chmin(id);
            path_range[path_root[v]].1.chmax(id + 1);
            left[v] = id;
            id += 1;
            dfs.push((v, true));
            let child = &child[v];
            if !child.is_empty() {
                for &u in child.iter().skip(1) {
                    path_root[u] = u;
                    dfs.push((u, false));
                }
                let u = child[0];
                path_root[u] = path_root[v];
                dfs.push((u, false));
            }
        }
        let mut inverse = vec![size; size];
        for (i, l) in left.iter().enumerate() {
            inverse[*l] = i;
        }
        self.child = child;
        self.parent = parent;
        self.left = left;
        self.right = right;
        self.path_root = path_root;
        self.inverse = inverse;
        self.path_range = path_range;
    }
    pub fn lca(&self, mut a: usize, mut b: usize) -> usize {
        assert!(a < self.size && b < self.size);
        let path = &self.path_root;
        let parent = &self.parent;
        let index = &self.left;
        while path[a] != path[b] {
            if index[a] > index[b] {
                a = parent[path[a]];
            } else {
                b = parent[path[b]];
            }
        }
        (index[a], a).min((index[b], b)).1
    }
    pub fn path_range(&self, v: usize) -> (usize, usize) {
        self.path_range[self.path_root[v]]
    }
    pub fn path(
        &self,
        src: usize,
        dst: usize,
        up: &mut Vec<(usize, usize)>,
        down: &mut Vec<(usize, usize)>,
    ) {
        assert!(src < self.size && dst < self.size);
        up.clear();
        down.clear();
        let path = &self.path_root;
        let parent = &self.parent;
        let index = &self.left;
        let mut x = src;
        let mut y = dst;
        while path[x] != path[y] {
            if index[x] > index[y] {
                let p = path[x];
                assert!(p == path[p]);
                up.push((index[p], index[x] + 1));
                x = parent[p];
            } else {
                let p = path[y];
                assert!(p == path[p]);
                down.push((index[p], index[y] + 1));
                y = parent[p];
            }
        }
        if index[x] <= index[y] {
            down.push((index[x], index[y] + 1));
        } else {
            up.push((index[y], index[x] + 1));
        }
        down.reverse();
    }
    pub fn sub_tree(&self, v: usize) -> (usize, usize) {
        assert!(v < self.size, "ERROR: {}", v);
        (self.left[v], self.right[v])
    }
    pub fn parent(&self, v: usize) -> Option<usize> {
        assert!(v < self.size);
        let p = self.parent[v];
        if p == v {
            None
        } else {
            Some(p)
        }
    }
    // s -> t へのパスの2番目の頂点を返す
    pub fn next(&self, s: usize, t: usize) -> usize {
        assert!(s < self.size && t < self.size && s != t);
        let (a, b) = self.sub_tree(s);
        let (c, d) = self.sub_tree(t);
        if !(a <= c && d <= b) {
            return self.parent[s];
        }
        let mut pos = t;
        let mut pre = t;
        while self.path_root[s] != self.path_root[pos] {
            pre = self.path_root[pos];
            pos = self.parent[pre];
        }
        if s == pos {
            pre
        } else {
            self.child[s][0]
        }
    }
    pub fn vertex(&self, x: usize) -> usize {
        assert!(x < self.size);
        self.inverse[x]
    }
}
// ---------- end Heavy-Light decomposition ----------
// ---------- begin SegmentTree Point update Range query ----------
mod segment_tree {
    pub struct PURQ<T, F> {
        size: usize,
        data: Vec<T>,
        e: T,
        op: F,
    }
    #[allow(dead_code)]
    impl<T, F> PURQ<T, F>
    where
        T: Clone,
        F: Fn(&T, &T) -> T,
    {
        pub fn new(size: usize, e: T, op: F) -> PURQ<T, F> {
            let size = size.next_power_of_two();
            PURQ {
                size,
                data: vec![e.clone(); 2 * size],
                e: e,
                op: op,
            }
        }
        pub fn update(&mut self, x: usize, v: T) {
            assert!(x < self.size);
            let mut x = x + self.size;
            let data = &mut self.data;
            data[x] = v;
            x >>= 1;
            while x > 0 {
                data[x] = (self.op)(&data[2 * x], &data[2 * x + 1]);
                x >>= 1;
            }
        }
        pub fn update_tmp(&mut self, x: usize, v: T) {
            assert!(x < self.size);
            self.data[x + self.size] = v;
        }
        pub fn update_all(&mut self) {
            let data = &mut self.data;
            for k in (1..self.size).rev() {
                data[k] = (self.op)(&data[2 * k], &data[2 * k + 1]);
            }
        }
        pub fn find(&self, l: usize, r: usize) -> T {
            assert!(l <= r && r <= self.size);
            if l == r {
                return self.e.clone();
            }
            let mut p = self.e.clone();
            let mut q = self.e.clone();
            let mut l = l + self.size;
            let mut r = r + self.size;
            let data = &self.data;
            while l < r {
                if l & 1 == 1 {
                    p = (self.op)(&p, &data[l]);
                    l += 1;
                }
                if r & 1 == 1 {
                    r -= 1;
                    q = (self.op)(&data[r], &q);
                }
                l >>= 1;
                r >>= 1;
            }
            (self.op)(&p, &q)
        }
    }
}
// ---------- end SegmentTree Point update Range query ----------
// ---------- begin input macro ----------
// reference: https://qiita.com/tanakh/items/0ba42c7ca36cd29d0ac8
macro_rules! input {
    (source = $s:expr, $($r:tt)*) => {
        let mut iter = $s.split_whitespace();
        input_inner!{iter, $($r)*}
    };
    ($($r:tt)*) => {
        let s = {
            use std::io::Read;
            let mut s = String::new();
            std::io::stdin().read_to_string(&mut s).unwrap();
            s
        };
        let mut iter = s.split_whitespace();
        input_inner!{iter, $($r)*}
    };
}

macro_rules! input_inner {
    ($iter:expr) => {};
    ($iter:expr, ) => {};
    ($iter:expr, $var:ident : $t:tt $($r:tt)*) => {
        let $var = read_value!($iter, $t);
        input_inner!{$iter $($r)*}
    };
}

macro_rules! read_value {
    ($iter:expr, ( $($t:tt),* )) => {
        ( $(read_value!($iter, $t)),* )
    };
    ($iter:expr, [ $t:tt ; $len:expr ]) => {
        (0..$len).map(|_| read_value!($iter, $t)).collect::<Vec<_>>()
    };
    ($iter:expr, chars) => {
        read_value!($iter, String).chars().collect::<Vec<char>>()
    };
    ($iter:expr, bytes) => {
        read_value!($iter, String).bytes().collect::<Vec<u8>>()
    };
    ($iter:expr, usize1) => {
        read_value!($iter, usize) - 1
    };
    ($iter:expr, $t:ty) => {
        $iter.next().unwrap().parse::<$t>().expect("Parse error")
    };
}
// ---------- end input macro ----------

fn run(e: Vec<(usize, usize)>, ask: Vec<(usize, usize)>) -> Vec<usize> {
    let n = e.len() + 1;
    let root = 0;
    let mut hld = HLD::new(n);
    let mut g = vec![vec![]; n];
    for (a, b) in e {
        hld.add_edge(a, b);
        g[a].push(b);
        g[b].push(a);
    }
    hld.build(root);
    let hld = hld;
    let mut bfs_range = vec![(0, 0); n];
    let mut topo = vec![root];
    let mut index = vec![0; n];
    for i in 0..n {
        let v = topo[i];
        index[v] = i;
        bfs_range[v] = (topo.len(), topo.len() + g[v].len());
        for u in g[v].clone() {
            g[u].retain(|e| *e != v);
            topo.push(u);
        }
    }
    let bfs_range_without_heavy = |v: usize| -> ((usize, usize), (usize, usize)) {
        let pos = hld.sub_tree(v).0;
        let range = hld.path_range(v);
        let (l, r) = bfs_range[v];
        if pos + 1 < range.1 {
            let u = hld.vertex(pos + 1);
            let x = index[u];
            ((l, x), (x + 1, r))
        } else {
            ((l, r), (0, 0))
        }
    };
    let mut value = (1..=n).collect::<Vec<_>>();
    let mut rmq = segment_tree::PURQ::new(n, (0, 0), |a, b| std::cmp::max(*a, *b));
    for (i, v) in value.iter().enumerate() {
        rmq.update_tmp(index[i], (*v, i));
    }
    rmq.update_all();
    let empty = n + 2;
    let mut down = segment_tree::PURQ::new(n, (empty, false), |a, b| {
        if a.0 == empty {
            *b
        } else if b.0 == empty {
            *a
        } else if a.1 {
            *b
        } else {
            *a
        }
    });
    for i in 0..n {
        let path_range = hld.path_range(i);
        let p = bfs_range[i];
        let max = rmq.find(p.0, p.1);
        let pos = hld.sub_tree(i).0;
        let go = pos + 1 < path_range.1 && value[hld.vertex(pos + 1)] == max.0;
        down.update_tmp(pos, (i, go));
    }
    down.update_all();
    let mut up = segment_tree::PURQ::new(n, (empty, false), |a, b| {
        if a.0 == empty {
            *b
        } else if b.0 == empty {
            *a
        } else if b.1 {
            *a
        } else {
            *b
        }
    });
    for i in 0..n {
        let path_range = hld.path_range(i);
        let (a, b) = bfs_range_without_heavy(i);
        let max = rmq.find(a.0, a.1).max(rmq.find(b.0, b.1));
        let pos = hld.sub_tree(i).0;
        let go = path_range.0 < pos && value[hld.vertex(pos - 1)] > max.0;
        up.update_tmp(pos, (i, go));
    }
    up.update_all();
    let mut ans = vec![0; ask.len()];
    let mut x = 0;
    for (it, (ans, (a, b))) in ans.iter_mut().zip(ask).enumerate() {
        let a = (a + x) % n;
        let b = (b + x) % n;
        value.swap(a, b);
        rmq.update(index[a], (value[a], a));
        rmq.update(index[b], (value[b], b));
        for &v in [a, b].iter() {
            if let Some(i) = hld.parent(v) {
                let path_range = hld.path_range(i);
                let p = bfs_range[i];
                let max = rmq.find(p.0, p.1);
                let pos = hld.sub_tree(i).0;
                let go = pos + 1 < path_range.1 && value[hld.vertex(pos + 1)] == max.0;
                down.update(pos, (i, go));
                let path_range = hld.path_range(i);
                let (a, b) = bfs_range_without_heavy(i);
                let max = rmq.find(a.0, a.1).max(rmq.find(b.0, b.1));
                let pos = hld.sub_tree(i).0;
                let go = path_range.0 < pos && value[hld.vertex(pos - 1)] > max.0;
                up.update(pos, (i, go));
            }
            let path_range = hld.path_range(v);
            let pos = hld.sub_tree(v).0;
            if pos + 1 < path_range.1 {
                let i = hld.vertex(pos + 1);
                let path_range = hld.path_range(i);
                let (a, b) = bfs_range_without_heavy(i);
                let max = rmq.find(a.0, a.1).max(rmq.find(b.0, b.1));
                let pos = hld.sub_tree(i).0;
                let go = path_range.0 < pos && value[hld.vertex(pos - 1)] > max.0;
                up.update(pos, (i, go));
            }
        }
        let mut now = a;
        let mut del = empty;
        loop {
            let (l, r) = bfs_range[now];
            let mut max = (0, 0);
            if index.get(del).map_or(false, |d| l <= *d && *d < r) {
                let del_pos = index[del];
                max.chmax(rmq.find(l, del_pos));
                max.chmax(rmq.find(del_pos + 1, r));
            } else {
                max = rmq.find(l, r);
            }
            let (l, r) = hld.path_range(now);
            let x = hld.sub_tree(now).0;
            if max.0 == 0 && (now == 0 || hld.parent(now).unwrap() == del) {
                break;
            }
            if hld
                .parent(now)
                .map_or(false, |p| p != del && value[p] > max.0)
            {
                if l == x {
                    del = now;
                    now = hld.parent(now).unwrap();
                } else {
                    let (p, y) = up.find(l, x);
                    assert!(!y);
                    now = p;
                    del = hld.vertex(hld.sub_tree(now).0 + 1);
                }
            } else if x + 1 < r && max.1 == hld.vertex(x + 1) {
                let (c, y) = down.find(x + 1, r);
                assert!(!y);
                assert!(c < n);
                now = c;
                del = hld.vertex(hld.sub_tree(now).0 - 1);
            } else {
                assert!(max.0 > 0);
                del = now;
                now = max.1;
            }
        }
        x = now + 1;
        *ans = x;
    }
    ans
}

fn naive(e: Vec<(usize, usize)>, ask: Vec<(usize, usize)>) -> Vec<usize> {
    let n = e.len() + 1;
    let mut g = vec![vec![]; n];
    for (a, b) in e {
        g[a].push(b);
        g[b].push(a);
    }
    let mut a = (0..n).collect::<Vec<_>>();
    let mut ans = vec![0; ask.len()];
    let mut x = 0;
    for (ans, (u, v)) in ans.iter_mut().zip(ask) {
        let u = (u + x) % n;
        let v = (v + x) % n;
        a.swap(u, v);
        let mut pre = n;
        let mut pos = u;
        while let Some(&v) = g[pos].iter().filter(|p| **p != pre).max_by_key(|v| a[**v]) {
            pre = pos;
            pos = v;
        }
        x = pos + 1;
        *ans = x;
    }
    ans
}

fn main() {
    input! {
        n: usize,
        e: [(usize1, usize1); n - 1],
        q: usize,
        ask: [(usize1, usize1); q],
    }
    let test = run(e.clone(), ask.clone());
    use std::io::Write;
    let out = std::io::stdout();
    let mut out = std::io::BufWriter::new(out.lock());
    for t in test {
        writeln!(out, "{}", t).ok();
    }
    /*
    let n = 1000;
    for _ in 0..1000 {
        let mut e = vec![];
        let mut dsu = DSU::new(n);
        while dsu.size(0) < n {
            let a = rand() % n;
            let b = rand() % n;
            if dsu.unite(a, b).is_some() {
                e.push((a, b));
            }
        }
        let ask = (0..n)
            .map(|_| {
                let mut a = rand() % n;
                let mut b = rand() % n;
                while a == b {
                    a = rand() % n;
                    b = rand() % n;
                }
                (a, b)
            })
            .collect::<Vec<_>>();
        let test = run(e.clone(), ask.clone());
        let ans = naive(e.clone(), ask.clone());
        if test != ans {
            println!("{:?}", e);
            println!("{:?}", ask);
            assert_eq!(test, ans);
        }
    }
    */
}

yukicoder

結果

コンパイルメッセージ

ソースコード