#252597 (Rust) No.650 行列木クエリ

• 問題: No.650 行列木クエリ
• ユーザー: lzy9
• 提出日時: 2018-04-20 00:09:47
• 言語: Rust (1.77.0)
• 結果: AC
• 実行時間: 156 ms / 2,000 ms
• コード長: 7,458 bytes
• コンパイル時間: 3,632 ms
• コンパイル使用メモリ: 160,528 KB
• 実行使用メモリ: 41,468 KB
• 最終ジャッジ日時: 2023-09-09 11:51:40
• 合計ジャッジ時間: 5,499 ms
• ジャッジサーバーID （参考情報）: judge11 / judge13

入力	結果	実行時間; 実行使用メモリ
testcase_00	AC	1 ms; 4,380 KB
testcase_01	AC	53 ms; 7,556 KB
testcase_02	AC	156 ms; 30,584 KB
testcase_03	AC	1 ms; 4,380 KB
testcase_04	AC	53 ms; 7,620 KB
testcase_05	AC	150 ms; 30,172 KB
testcase_06	AC	1 ms; 4,376 KB
testcase_07	AC	1 ms; 4,376 KB
testcase_08	AC	58 ms; 10,212 KB
testcase_09	AC	120 ms; 41,468 KB
testcase_10	AC	1 ms; 4,376 KB

ソースコード

#[allow(unused_imports)]
use std::io::*;
#[allow(unused_imports)]
use std::str::FromStr;
#[allow(unused_imports)]
use std::cmp::{min, max};
#[allow(unused_imports)]
use std::mem::swap;
#[allow(unused_imports)]
use std::collections::{HashMap, VecDeque};

#[allow(dead_code)]
fn read<T: FromStr>() -> T {
    let stdin = stdin();
    let stdin_lock = stdin.lock();
    let s = stdin_lock
        .bytes()
        .map(|c| c.unwrap() as char)
        .skip_while(|c| c.is_whitespace())
        .take_while(|c| !c.is_whitespace())
        .collect::<String>();
    s.parse::<T>().ok().unwrap()
}

#[allow(dead_code)]
static DX: &'static [i32] = &[0, 0, 1, -1];
#[allow(dead_code)]
static DY: &'static [i32] = &[1, -1, 0, 0];
#[allow(dead_code)]
static MOD: u64 = 1000000007;

fn main() {
    let n: usize = read();
    let mut graph: Vec<Vec<usize>> = (0..n).map(|_| Vec::new()).collect();
    let mut edge: Vec<[usize; 2]> = Vec::new();

    for _ in 0..n - 1 {
        let a: usize = read();
        let b: usize = read();

        graph[a].push(b);
        graph[b].push(a);

        edge.push([a, b]);
    }

    let (parent, depth, chain, index, head) = HLD::make(&graph, 0);
    let mut segtree: Vec<SegTree<[u64; 4]>> = chain.iter().map(|c| SegTree::new(c.len(), mul, [1, 0, 0, 1])).collect();

    let q: usize = read();

    for _ in 0..q {
        let t: char = read();

        if t == 'x' {
            let mut i: usize = read();
            let x: [u64; 4] = [read(), read(), read(), read()];

            if depth[edge[i][0]] < depth[edge[i][1]] {
                i = edge[i][1];
            } else {
                i = edge[i][0];
            }

            let (cidx, idx) = index[i];
            segtree[cidx].update(idx, x);
        } else {
            let from: usize = read();
            let mut to: usize = read();
            let mut ans: [u64; 4] = [1, 0, 0, 1];

            while index[from].0 != index[to].0 {
                let (cidx, idx) = index[to];

                ans = mul(segtree[cidx].query(0, idx + 1), ans);
                to = parent[head[to]].unwrap();
            }

            let (cidx, fidx) = index[from];
            let (_, tidx) = index[to];

            ans = mul(segtree[cidx].query(fidx + 1, tidx + 1), ans);

            for i in 0..3 {
                print!("{} ", ans[i]);
            }

            println!("{}", ans[3]);
        }
    }
}

fn mul(a: [u64; 4], b: [u64; 4]) -> [u64; 4] {
    let mut ret = [0 as u64; 4];

    ret[0] = ((a[0] * b[0] % MOD) + (a[1] * b[2] % MOD)) % MOD;
    ret[1] = ((a[0] * b[1] % MOD) + (a[1] * b[3] % MOD)) % MOD;
    ret[2] = ((a[2] * b[0] % MOD) + (a[3] * b[2] % MOD)) % MOD;
    ret[3] = ((a[2] * b[1] % MOD) + (a[3] * b[3] % MOD)) % MOD;

    ret
}

#[allow(dead_code)]
struct HLD {
    parent: Vec<Option<usize>>,
    depth: Vec<usize>,
    chain: Vec<Vec<usize>>,
    index: Vec<(usize, usize)>,
    head: Vec<usize>
}

impl HLD {
    #[allow(dead_code)]
    fn new(graph: &Vec<Vec<usize>>, root: usize) -> Self {
        let (parent, depth, chain, index, head) = HLD::make(graph, root);

        HLD {
            parent,
            depth,
            chain,
            index,
            head,
        }
    }

    #[allow(dead_code)]
    fn make(graph: &Vec<Vec<usize>>, root: usize) -> (Vec<Option<usize>>, Vec<usize>, Vec<Vec<usize>>, Vec<(usize, usize)>, Vec<usize>) {
        let n = graph.len();
        let mut size = vec![0; n];
        let mut parent: Vec<Option<usize>> = vec![None; n];
        let mut depth: Vec<usize> = vec![0; n];
        let mut heavy: Vec<Option<usize>> = vec![None; n];

        HLD::make_dfs(graph, root, &mut size, &mut parent, &mut depth, &mut heavy);
        let (chain, index, head) = HLD::make_chain(n, graph, &parent, &heavy, root);

        (parent, depth, chain, index, head)
    }

    #[allow(dead_code)]
    fn make_chain(n: usize, graph: &Vec<Vec<usize>>, parent: &Vec<Option<usize>>, heavy: &Vec<Option<usize>>, root: usize) -> (Vec<Vec<usize>>, Vec<(usize, usize)>, Vec<usize>) {
        let mut chain: Vec<Vec<usize>> = Vec::new();
        let mut index: Vec<(usize, usize)> = vec![(0, 0); n];
        let mut head: Vec<usize> = vec![0; n];
        let mut queued = vec![false; n];
        let mut idx: usize = 0;

        let mut queue = VecDeque::new();
        queue.push_back(root);
        queued[root] = true;

        while let Some(node) = queue.pop_front() {
            if parent[node] == None || heavy[parent[node].unwrap()] != Some(node) {
                chain.push(Vec::new());
                chain[idx].push(node);
                head[node] = node;
                index[node] = (idx, 0);
                idx += 1;
            } else {
                let (cidx, idx) = index[parent[node].unwrap()];
                chain[cidx].push(node);
                index[node] = (cidx, idx + 1);
                head[node] = head[parent[node].unwrap()];
            }

            for child in &graph[node] {
                if queued[*child] {
                    continue;
                }

                queue.push_back(*child);
                queued[*child] = true;
            }
        }

        (chain, index, head)
    }

    #[allow(dead_code)]
    fn make_dfs(
        graph: &Vec<Vec<usize>>,
        current: usize,
        size: &mut Vec<usize>,
        parent: &mut Vec<Option<usize>>,
        depth: &mut Vec<usize>,
        heavy: &mut Vec<Option<usize>>)
    {
        size[current] = 1;
        let mut heaviest: Option<usize> = None;

        for child in graph[current].iter() {
            if let Some(par) = parent[current] {
                if par == *child {
                    continue;
                }
            }

            parent[*child] = Some(current);
            depth[*child] = depth[current] + 1;
            HLD::make_dfs(graph, *child, size, parent, depth, heavy);

            if heaviest == None || size[heaviest.unwrap()] < size[*child] {
                heaviest = Some(*child);
            }

            size[current] += size[*child];
        }

        heavy[current] = heaviest;
    }
}

#[allow(dead_code)]
struct SegTree<T> where T: Clone + Copy {
    n: usize,
    dat: Vec<T>,
    operation: fn(T, T) -> T,
    default: T,
}

impl<T> SegTree<T> where T: Clone + Copy {
    #[allow(dead_code)]
    fn new(n: usize, operation: fn(T, T) -> T, default: T) -> Self {
        let mut size = 1;

        while size < n {
            size <<= 1;
        }

        SegTree {
            n: size,
            dat: vec![default; size * 2],
            operation,
            default,
        }
    }

    #[allow(dead_code)]
    fn update(&mut self, idx: usize, x: T) {
        let mut k = idx + self.n - 1;

        self.dat[k] = x;

        while 0 < k {
            k = (k - 1) / 2;

            self.dat[k] = (self.operation)(self.dat[k * 2 + 1], self.dat[k * 2 + 2]);
        }
    }

    #[allow(dead_code)]
    fn query(&self, from: usize, to: usize) -> T {
        self.query_rec(from, to, 0, 0, self.n)
    }

    #[allow(dead_code)]
    fn query_rec(&self, from: usize, to: usize, idx: usize, a: usize, b: usize) -> T {
        if b <= from || to <= a {
            return self.default;
        }

        if from <= a && b <= to {
            return self.dat[idx];
        }

        let mid = (a + b) / 2;

        (self.operation)(self.query_rec(from, to, idx * 2 + 1, a, mid), self.query_rec(from, to, idx * 2 + 2, mid, b))
    }
}