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

• 問題: No.650 行列木クエリ
• ユーザー: hatoo
• 提出日時: 2018-02-15 18:04:45
• 言語: Rust (1.77.0)
• 結果: AC
• 実行時間: 105 ms / 2,000 ms
• コード長: 10,189 bytes
• コンパイル時間: 2,419 ms
• コンパイル使用メモリ: 164,396 KB
• 実行使用メモリ: 35,644 KB
• 最終ジャッジ日時: 2023-08-27 02:29:34
• 合計ジャッジ時間: 4,279 ms
• ジャッジサーバーID （参考情報）: judge15 / judge14

入力	結果	実行時間; 実行使用メモリ
testcase_00	AC	1 ms; 4,380 KB
testcase_01	AC	38 ms; 7,816 KB
testcase_02	AC	101 ms; 31,600 KB
testcase_03	AC	1 ms; 4,376 KB
testcase_04	AC	36 ms; 7,852 KB
testcase_05	AC	105 ms; 31,412 KB
testcase_06	AC	1 ms; 4,380 KB
testcase_07	AC	1 ms; 4,380 KB
testcase_08	AC	37 ms; 8,592 KB
testcase_09	AC	83 ms; 35,644 KB
testcase_10	AC	1 ms; 4,376 KB

ソースコード

/**
 *  _           _                 __                            _   _ _   _                                 _                    _                  _
 * | |         | |               / /                           | | (_) | (_)                               | |                  (_)                | |
 * | |__   __ _| |_ ___   ___   / /__ ___  _ __ ___  _ __   ___| |_ _| |_ ___   _____ ______ _ __ _   _ ___| |_ ______ ___ _ __  _ _ __  _ __   ___| |_ ___
 * | '_ \ / _` | __/ _ \ / _ \ / / __/ _ \| '_ ` _ \| '_ \ / _ \ __| | __| \ \ / / _ \______| '__| | | / __| __|______/ __| '_ \| | '_ \| '_ \ / _ \ __/ __|
 * | | | | (_| | || (_) | (_) / / (_| (_) | | | | | | |_) |  __/ |_| | |_| |\ V /  __/      | |  | |_| \__ \ |_       \__ \ | | | | |_) | |_) |  __/ |_\__ \
 * |_| |_|\__,_|\__\___/ \___/_/ \___\___/|_| |_| |_| .__/ \___|\__|_|\__|_| \_/ \___|      |_|   \__,_|___/\__|      |___/_| |_|_| .__/| .__/ \___|\__|___/
 *                                                  | |                                                                           | |   | |
 *                                                  |_|                                                                           |_|   |_|
 *
 * https://github.com/hatoo/competitive-rust-snippets
 */
#[allow(unused_imports)]
use std::cmp::{max, min, Ordering};
#[allow(unused_imports)]
use std::collections::{BTreeMap, BTreeSet, BinaryHeap, HashMap, HashSet, VecDeque};
#[allow(unused_imports)]
use std::iter::FromIterator;
#[allow(unused_imports)]
use std::io::{stdin, stdout, BufWriter, Write};
mod util {
    use std::io::{stdin, stdout, BufWriter, StdoutLock};
    use std::str::FromStr;
    use std::fmt::Debug;
    #[allow(dead_code)]
    pub fn line() -> String {
        let mut line: String = String::new();
        stdin().read_line(&mut line).unwrap();
        line.trim().to_string()
    }
    #[allow(dead_code)]
    pub fn gets<T: FromStr>() -> Vec<T>
    where
        <T as FromStr>::Err: Debug,
    {
        let mut line: String = String::new();
        stdin().read_line(&mut line).unwrap();
        line.split_whitespace()
            .map(|t| t.parse().unwrap())
            .collect()
    }
    #[allow(dead_code)]
    pub fn with_bufwriter<F: FnOnce(BufWriter<StdoutLock>) -> ()>(f: F) {
        let out = stdout();
        let writer = BufWriter::new(out.lock());
        f(writer)
    }
}
#[allow(unused_macros)]
macro_rules ! get { ( $ t : ty ) => { { let mut line : String = String :: new ( ) ; stdin ( ) . read_line ( & mut line ) . unwrap ( ) ; line . trim ( ) . parse ::<$ t > ( ) . unwrap ( ) } } ; ( $ ( $ t : ty ) ,* ) => { { let mut line : String = String :: new ( ) ; stdin ( ) . read_line ( & mut line ) . unwrap ( ) ; let mut iter = line . split_whitespace ( ) ; ( $ ( iter . next ( ) . unwrap ( ) . parse ::<$ t > ( ) . unwrap ( ) , ) * ) } } ; ( $ t : ty ; $ n : expr ) => { ( 0 ..$ n ) . map ( | _ | get ! ( $ t ) ) . collect ::< Vec < _ >> ( ) } ; ( $ ( $ t : ty ) ,*; $ n : expr ) => { ( 0 ..$ n ) . map ( | _ | get ! ( $ ( $ t ) ,* ) ) . collect ::< Vec < _ >> ( ) } ; ( $ t : ty ;; ) => { { let mut line : String = String :: new ( ) ; stdin ( ) . read_line ( & mut line ) . unwrap ( ) ; line . split_whitespace ( ) . map ( | t | t . parse ::<$ t > ( ) . unwrap ( ) ) . collect ::< Vec < _ >> ( ) } } ; }
#[allow(unused_macros)]
macro_rules ! debug { ( $ ( $ a : expr ) ,* ) => { println ! ( concat ! ( $ ( stringify ! ( $ a ) , " = {:?}, " ) ,* ) , $ ( $ a ) ,* ) ; } }

#[allow(dead_code)]
pub trait Monoid {
    type T: Clone;
    fn id() -> Self::T;
    fn op(a: &Self::T, b: &Self::T) -> Self::T;
}
#[allow(dead_code)]
/// Segment Tree
pub struct SEG<M: Monoid> {
    n: usize,
    buf: Vec<M::T>,
}
impl<M: Monoid> SEG<M> {
    #[allow(dead_code)]
    pub fn new(n: usize) -> SEG<M> {
        SEG {
            n: n,
            buf: vec![M::id().clone(); 2 * n],
        }
    }
    #[allow(dead_code)]
    pub fn update(&mut self, k: usize, a: M::T) {
        let mut k = k + self.n;
        self.buf[k] = a;
        while k > 0 {
            k >>= 1;
            self.buf[k] = M::op(&self.buf[k << 1], &self.buf[(k << 1) | 1]);
        }
    }
    #[allow(dead_code)]
    pub fn add(&mut self, k: usize, a: &M::T) {
        let mut k = k + self.n;
        self.buf[k] = M::op(&self.buf[k], a);
        while k > 0 {
            k >>= 1;
            self.buf[k] = M::op(&self.buf[k << 1], &self.buf[(k << 1) | 1]);
        }
    }
    #[allow(dead_code)]
    fn query(&self, l: usize, r: usize) -> Option<M::T> {
        let combine = |resl, resr| match (resl, resr) {
            (Some(l), Some(r)) => Some(M::op(&l, &r)),
            (Some(l), None) => Some(l),
            (None, Some(r)) => Some(r),
            _ => None,
        };
        let mut vl = None;
        let mut vr = None;
        let mut l = l + self.n;
        let mut r = r + self.n;
        while l < r {
            if l & 1 == 1 {
                vl = combine(vl, Some(self.buf[l].clone()));
                l += 1;
            }
            if r & 1 == 1 {
                r -= 1;
                vr = combine(Some(self.buf[r].clone()), vr);
            }
            l >>= 1;
            r >>= 1;
        }
        combine(vl, vr)
    }
}

struct Tree {
    root: usize,
    parent: Vec<Option<usize>>,
    childs: Vec<Vec<usize>>,
}

impl Tree {
    pub fn from_neighbor_list(n: usize, root: usize, g: &[Vec<usize>]) -> Tree {
        let mut parent = vec![None; n];
        let mut childs = vec![Vec::new(); n];

        let mut stack = vec![(root, None)];

        while let Some((i, p)) = stack.pop() {
            parent[i] = p;

            for &to in &g[i] {
                if Some(to) != p {
                    stack.push((to, Some(i)));
                    childs[i].push(to);
                }
            }
        }

        Tree {
            root: root,
            parent: parent,
            childs: childs,
        }
    }
}

struct HeavyLightDecomposition {
    ids: Vec<(usize, usize)>,
    parents: Vec<Option<(usize, usize)>>,
    parts: Vec<Vec<usize>>,
}

impl HeavyLightDecomposition {
    fn new(tree: &Tree) -> HeavyLightDecomposition {
        fn size(i: usize, tree: &Tree, memo: &mut [Option<usize>]) -> usize {
            if let Some(res) = memo[i] {
                return res;
            }

            let res = tree.childs[i]
                .iter()
                .map(|&to| size(to, tree, memo))
                .sum::<usize>() + 1;
            memo[i] = Some(res);
            res
        }

        let n = tree.parent.len();
        let root = tree.root;
        let mut memo = vec![None; n];

        let mut ids = vec![(0, 0); n];
        let mut parts: Vec<Vec<usize>> = Vec::new();
        let mut stack = vec![(root, false, None)];
        let mut parents = Vec::new();

        while let Some((i, h, pid)) = stack.pop() {
            if h {
                let (k, _) = pid.unwrap();
                ids[i] = (k, parts[k].len());
                parts[k].push(i);
            } else {
                ids[i] = (parts.len(), 0);
                parts.push(vec![i]);
                parents.push(pid);
            }

            let id = ids[i];

            let heavy = tree.childs[i]
                .iter()
                .max_by_key(|&&to| size(to, &tree, &mut memo));

            for &to in &tree.childs[i] {
                if Some(&to) != heavy {
                    stack.push((to, false, Some(id)));
                }
            }

            if let Some(&h) = heavy {
                stack.push((h, true, Some(id)));
            }
        }

        HeavyLightDecomposition {
            ids: ids,
            parents: parents,
            parts: parts,
        }
    }
}

#[allow(dead_code)]
pub const M: u64 = 1_000_000_007;

struct MatMul;
impl Monoid for MatMul {
    type T = [u64; 4];

    fn id() -> Self::T {
        [1, 0, 0, 1]
    }

    fn op(a: &Self::T, b: &Self::T) -> Self::T {
        [
            (a[0] * b[0] % M + a[1] * b[2] % M) % M,
            (a[0] * b[1] % M + a[1] * b[3] % M) % M,
            (a[2] * b[0] % M + a[3] * b[2] % M) % M,
            (a[2] * b[1] % M + a[3] * b[3] % M) % M,
        ]
    }
}

#[allow(dead_code)]
fn main() {
    let n = get!(usize);
    let ab = get!(usize, usize; n-1);
    let q = get!(usize);

    let mut g = vec![Vec::new(); n];

    for &(a, b) in &ab {
        g[a].push(b);
        g[b].push(a);
    }

    let tree = Tree::from_neighbor_list(n, 0, &g);
    let hld = HeavyLightDecomposition::new(&tree);
    let mut segs: Vec<SEG<MatMul>> = hld.parts.iter().map(|p| SEG::new(p.len())).collect();

    for _ in 0..q {
        let input = util::line();
        let mut split = input.split_whitespace();

        if split.next() == Some("x") {
            let i: usize = split.next().unwrap().parse().unwrap();

            let mat = [
                split.next().unwrap().parse().unwrap(),
                split.next().unwrap().parse().unwrap(),
                split.next().unwrap().parse().unwrap(),
                split.next().unwrap().parse().unwrap(),
            ];

            let (a, b) = ab[i];

            let i = if tree.parent[a] == Some(b) { a } else { b };

            let (p, k) = hld.ids[i];

            segs[p].update(k, mat);
        } else {
            let i: usize = split.next().unwrap().parse().unwrap();
            let j: usize = split.next().unwrap().parse().unwrap();

            let (mut p, mut k) = hld.ids[j];
            let mut mat = MatMul::id();

            while p != hld.ids[i].0 {
                let m1 = segs[p].query(0, k + 1).unwrap_or_else(|| MatMul::id());
                mat = MatMul::op(&m1, &mat);
                let (p1, k1) = hld.parents[p].unwrap();
                p = p1;
                k = k1;
            }
            let m1 = segs[p]
                .query(hld.ids[i].1 + 1, k + 1)
                .unwrap_or_else(|| MatMul::id());
            mat = MatMul::op(&m1, &mat);

            println!(
                "{}",
                mat.iter()
                    .map(|x| x.to_string())
                    .collect::<Vec<_>>()
                    .join(" ")
            );
        }
    }
}