library_checker/number_theory/
stern_brocot_tree.rs

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

library_checker/number_theory/
stern_brocot_tree.rs
