結果

問題 No.650 行列木クエリ
ユーザー lzy9lzy9
提出日時 2018-04-20 00:09:47
言語 Rust
(1.77.0 + proconio)
結果
AC  
実行時間 163 ms / 2,000 ms
コード長 7,458 bytes
コンパイル時間 28,911 ms
コンパイル使用メモリ 376,380 KB
実行使用メモリ 40,000 KB
最終ジャッジ日時 2024-06-27 04:57:03
合計ジャッジ時間 18,046 ms
ジャッジサーバーID
（参考情報） 		judge5 / judge4
このコードへのチャレンジ
（要ログイン）

テストケース

テストケース表示
入力 結果 実行時間
実行使用メモリ
testcase_00 AC 1 ms
5,248 KB
testcase_01 AC 47 ms
7,612 KB
testcase_02 AC 163 ms
30,480 KB
testcase_03 AC 1 ms
5,376 KB
testcase_04 AC 47 ms
7,576 KB
testcase_05 AC 153 ms
30,300 KB
testcase_06 AC 1 ms
5,376 KB
testcase_07 AC 1 ms
5,376 KB
testcase_08 AC 50 ms
10,012 KB
testcase_09 AC 103 ms
40,000 KB
testcase_10 AC 1 ms
5,376 KB
権限があれば一括ダウンロードができます

ソースコード

diff # 

#[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))
    }
}
0