library_checker/number_theory/
stern_brocot_tree.rs

1use competitive::prelude::*;
2#[doc(no_inline)]
3pub use competitive::{
4    algorithm::{SbtNode, SbtPath, SternBrocotTree},
5    num::URational,
6};
7
8competitive::define_enum_scan! {
9    enum Query: raw {
10        "ENCODE_PATH" => EncodePath { a: u32, b: u32 }
11        "DECODE_PATH" => DecodePath { k: usize, path: [(char, u32); k] }
12        "LCA" => Lca { a: u32, b: u32, c: u32, d: u32 }
13        "ANCESTOR" => Ancestor { k: u32, a: u32, b: u32 }
14        "RANGE" => Range { a: u32, b: u32 }
15    }
16}
17
18#[verify::library_checker("stern_brocot_tree")]
19pub fn stern_brocot_tree(reader: impl Read, mut writer: impl Write) {
20    let s = read_all_unchecked(reader);
21    let mut scanner = Scanner::new(&s);
22    scan!(scanner, t);
23    for _ in 0..t {
24        scan!(scanner, query: Query);
25        match query {
26            Query::EncodePath { a, b } => {
27                let path = SbtPath::from(URational::new(a, b));
28                let len = if path.path.first() == Some(&0) {
29                    path.path.len() - 1
30                } else {
31                    path.path.len()
32                };
33                write!(writer, "{}", len).ok();
34                for (i, count) in path.into_iter().enumerate() {
35                    if count == 0 {
36                        continue;
37                    }
38                    if i % 2 == 0 {
39                        write!(writer, " R {}", count).ok();
40                    } else {
41                        write!(writer, " L {}", count).ok();
42                    }
43                }
44                writeln!(writer).ok();
45            }
46            Query::DecodePath { path, .. } => {
47                let node: SbtNode<u32> = if path.first().is_some_and(|t| t.0 == 'L') {
48                    [0].into_iter()
49                        .chain(path.into_iter().map(|(_, c)| c))
50                        .collect()
51                } else {
52                    path.into_iter().map(|(_, c)| c).collect()
53                };
54                let val = node.eval();
55                writeln!(writer, "{} {}", val.num, val.den).ok();
56            }
57            Query::Lca { a, b, c, d } => {
58                let path1 = SbtPath::from(URational::new(a, b));
59                let path2 = SbtPath::from(URational::new(c, d));
60                let val = SbtNode::lca(path1, path2).eval();
61                writeln!(writer, "{} {}", val.num, val.den).ok();
62            }
63            Query::Ancestor { k, a, b } => {
64                let mut path = SbtPath::from(URational::new(a, b));
65                let depth = path.depth();
66                if k <= depth {
67                    path.up(depth - k);
68                    let val = path.eval();
69                    writeln!(writer, "{} {}", val.num, val.den).ok();
70                } else {
71                    writeln!(writer, "-1").ok();
72                }
73            }
74            Query::Range { a, b } => {
75                let node = SbtPath::from(URational::new(a, b)).to_node();
76                writeln!(
77                    writer,
78                    "{} {} {} {}",
79                    node.l.num, node.l.den, node.r.num, node.r.den
80                )
81                .ok();
82            }
83        }
84    }
85}
library_checker/number_theory/stern_brocot_tree.rs

library_checker/number_theory/
stern_brocot_tree.rs
