Skip to content

Commit d524b0f

Browse files
lrvideckislrvideckis
authored
Kth on path and related (#55)
* initial draft; it AC's on LCA * fmt * add some un-tested fns * fmt * add another lca test * refactor * forgot this * progress * add another test * golf * another test * add another test * fmt * nit * clippy * test subtree range with edge values * fmt * clippy * finished refactor * still has bugs * saving progress * nits * now AC's!!!!!!!!!!! * now ACs on yosupo * fix cargo test * finish docs * now works with capturing vars * golf * fix test * reorder * fix * currently WA, but saving progress * now ACs * fix clippy * fix clippy * golf * nit * these are the functions I wanna add * remove untested function for now * add docs * fix example * fix bug * nit * fix code cov * ok actually fixed code cov this time * fix code cov * simplify test cuz I found a better test * Revert "simplify test cuz I found a better test" This reverts commit 4ab8679. * saving progress * FIX CODE COV FOR THE LAST TIME!!!!!!!! * adding progress * fininshed * actually simplify test * adding to save progress * cleaned * golf name * refactor * fix clippy * I like this order of functions * finished docs * better constant factor in theory * saving progress * fmt * nit * fix test * revert * fix clippy * Squashed commit of the following: commit f837403 Author: Luke Videckis <lukevideckis@gmail.com> Date: Mon Jun 3 16:02:37 2024 -0700 adding progress commit 040b129 Merge: a79bad6 8631c75 Author: Luke Videckis <lukevideckis@gmail.com> Date: Mon Jun 3 15:51:11 2024 -0700 Merge branch 'kth_par_and_related' into no_depth_ary commit a79bad6 Author: Luke Videckis <lukevideckis@gmail.com> Date: Mon Jun 3 15:48:13 2024 -0700 saving progress * change interface * fix * now ACs * finish test * now ACs * fix TLE * nits * another speedup * revert * split apart test to avoid TLE * nit --------- Co-authored-by: Luke Videckis <lukevideckis@gmail.com>
1 parent 6f26708 commit d524b0f

File tree

5 files changed

+193
-9
lines changed

5 files changed

+193
-9
lines changed

‎Cargo.toml‎

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -125,3 +125,11 @@ path = "examples/graphs/hld_path_sum2_aizu.rs"
125125
[[example]]
126126
name = "hld_path_composite_yosupo"
127127
path = "examples/graphs/hld_path_composite_yosupo.rs"
128+
129+
[[example]]
130+
name = "hld_jump_on_tree_nodes"
131+
path = "examples/graphs/hld_jump_on_tree_nodes.rs"
132+
133+
[[example]]
134+
name = "hld_jump_on_tree_edges"
135+
path = "examples/graphs/hld_jump_on_tree_edges.rs"

‎examples/graphs/hld_jump_on_tree_edges.rs‎

Lines changed: 46 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,46 @@
1+
// verification-helper: PROBLEM https://judge.yosupo.jp/problem/jump_on_tree
2+
3+
use proconio::input;
4+
use programming_team_code_rust::graphs::hld::HLD;
5+
6+
fn main() {
7+
input! {
8+
n: usize,
9+
q: usize,
10+
}
11+
12+
let mut adj = vec![vec![]; n];
13+
for _ in 1..n {
14+
input! {
15+
u: usize,
16+
v: usize
17+
}
18+
adj[u].push(v);
19+
adj[v].push(u);
20+
}
21+
22+
let hld = HLD::new(&mut adj, true);
23+
24+
for _ in 0..q {
25+
input! {
26+
u: usize,
27+
v: usize,
28+
k: usize,
29+
}
30+
31+
match hld.kth_on_path(u, v, k) {
32+
Some(w) => {
33+
assert!(k <= hld.dist(u, v));
34+
assert!(hld.on_path(u, v, w));
35+
if w != u {
36+
assert!(!hld.on_path(v, w, u));
37+
}
38+
println!("{}", w);
39+
}
40+
None => {
41+
assert!(k > hld.dist(u, v));
42+
println!("-1");
43+
}
44+
}
45+
}
46+
}

‎examples/graphs/hld_jump_on_tree_nodes.rs‎

Lines changed: 59 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,59 @@
1+
// verification-helper: PROBLEM https://judge.yosupo.jp/problem/jump_on_tree
2+
3+
use proconio::input;
4+
use programming_team_code_rust::graphs::hld::HLD;
5+
6+
fn main() {
7+
input! {
8+
n: usize,
9+
q: usize,
10+
}
11+
12+
let mut adj = vec![vec![]; n];
13+
for _ in 1..n {
14+
input! {
15+
u: usize,
16+
v: usize
17+
}
18+
adj[u].push(v);
19+
adj[v].push(u);
20+
}
21+
22+
let hld = HLD::new(&mut adj, false);
23+
24+
fn dfs(u: usize, adj: &[Vec<usize>], d: &mut [usize]) {
25+
for &v in &adj[u] {
26+
d[v] = 1 + d[u];
27+
dfs(v, adj, d);
28+
}
29+
}
30+
31+
let mut d = vec![0; n];
32+
dfs(0, &adj, &mut d);
33+
34+
for _ in 0..q {
35+
input! {
36+
u: usize,
37+
v: usize,
38+
k: usize,
39+
}
40+
41+
let in_sub_naive = hld.kth_par(v, d[u].abs_diff(d[v])) == Some(u);
42+
assert_eq!(in_sub_naive, hld.in_sub(u, v));
43+
44+
match hld.kth_on_path(u, v, k) {
45+
Some(w) => {
46+
assert!(k < hld.dist(u, v));
47+
assert!(hld.on_path(u, v, w));
48+
if w != u {
49+
assert!(!hld.on_path(v, w, u));
50+
}
51+
println!("{}", w);
52+
}
53+
None => {
54+
assert!(k >= hld.dist(u, v));
55+
println!("-1");
56+
}
57+
}
58+
}
59+
}

‎examples/graphs/hld_path_composite_yosupo.rs‎

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -70,11 +70,11 @@ fn main() {
7070
x: usize
7171
}
7272
let (mut u_anc_val, mut v_anc_val) = (st_forwards.unit, st_backwards.unit);
73-
hld.path(u, v, |range, u_anc| {
74-
if u_anc {
75-
u_anc_val = (st_forwards.op)(&u_anc_val, &st_backwards.query(range));
76-
} else {
73+
hld.path(u, v, |range, v_anc| {
74+
if v_anc {
7775
v_anc_val = (st_forwards.op)(&st_forwards.query(range), &v_anc_val);
76+
} else {
77+
u_anc_val = (st_forwards.op)(&u_anc_val, &st_backwards.query(range));
7878
}
7979
});
8080
let res = (st_forwards.op)(&u_anc_val, &v_anc_val);

‎src/graphs/hld.rs‎

Lines changed: 76 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -30,6 +30,7 @@ pub struct HLD {
3030
pub p: Vec<Option<usize>>,
3131
/// time in
3232
pub tin: Vec<usize>,
33+
ord: Vec<usize>,
3334
siz: Vec<usize>,
3435
head: Vec<usize>,
3536
vals_edges: bool,
@@ -62,7 +63,8 @@ impl HLD {
6263
}
6364
let mut tin = vec![0; n];
6465
let mut head = vec![0; n];
65-
for (i, &u) in get_dfs_preorder(adj).iter().enumerate() {
66+
let ord = get_dfs_preorder(adj);
67+
for (i, &u) in ord.iter().enumerate() {
6668
tin[u] = i;
6769
for &v in &adj[u] {
6870
head[v] = if v == adj[u][0] { head[u] } else { v };
@@ -71,6 +73,7 @@ impl HLD {
7173
HLD {
7274
p,
7375
siz,
76+
ord,
7477
tin,
7578
head,
7679
vals_edges,
@@ -83,21 +86,21 @@ impl HLD {
8386
/// - Time: O(log n) calls to `f`
8487
/// - Space: O(1)
8588
pub fn path(&self, mut u: usize, mut v: usize, mut f: impl FnMut(Range<usize>, bool)) {
86-
let mut u_anc = false;
89+
let mut v_anc = true;
8790
loop {
8891
if self.tin[u] > self.tin[v] {
8992
std::mem::swap(&mut u, &mut v);
90-
u_anc = !u_anc;
93+
v_anc = !v_anc;
9194
}
9295
if self.head[u] == self.head[v] {
9396
break;
9497
}
95-
f(self.tin[self.head[v]]..self.tin[v] + 1, u_anc);
98+
f(self.tin[self.head[v]]..self.tin[v] + 1, v_anc);
9699
v = self.p[self.head[v]].unwrap();
97100
}
98101
f(
99102
self.tin[u] + self.vals_edges as usize..self.tin[v] + 1,
100-
u_anc,
103+
v_anc,
101104
);
102105
}
103106

@@ -126,4 +129,72 @@ impl HLD {
126129
v = self.p[self.head[v]].unwrap();
127130
}
128131
}
132+
133+
/// If !vals_edges, then gets number of nodes on path from u to v
134+
/// If vals_edges, then gets number of edges on path from u to v
135+
///
136+
/// # Complexity
137+
/// - Time: O(log n)
138+
/// - Space: O(1)
139+
pub fn dist(&self, u: usize, v: usize) -> usize {
140+
let mut dst = 0;
141+
self.path(u, v, |range, _| dst += range.len());
142+
dst
143+
}
144+
145+
/// Returns true iff v is in u's subtree
146+
///
147+
/// # Complexity
148+
/// - Time: O(1)
149+
/// - Space: O(1)
150+
pub fn in_sub(&self, u: usize, v: usize) -> bool {
151+
u == v || self.sub_tree(u).contains(&self.tin[v])
152+
}
153+
154+
/// Returns true iff w is on the path from u to v
155+
///
156+
/// # Complexity
157+
/// - Time: O(log n)
158+
/// - Space: O(1)
159+
pub fn on_path(&self, u: usize, v: usize, w: usize) -> bool {
160+
self.in_sub(self.lca(u, v), w) && (self.in_sub(w, u) || self.in_sub(w, v))
161+
}
162+
163+
/// Returns a node k edges up from u, or None
164+
///
165+
/// # Complexity
166+
/// - Time: O(log n)
167+
/// - Space: O(1)
168+
pub fn kth_par(&self, mut u: usize, mut k: usize) -> Option<usize> {
169+
loop {
170+
let len_path = self.tin[u] - self.tin[self.head[u]];
171+
if k <= len_path {
172+
return Some(self.ord[self.tin[u] - k]);
173+
}
174+
match self.p[self.head[u]] {
175+
Some(v) => u = v,
176+
None => return None,
177+
}
178+
k -= len_path + 1;
179+
}
180+
}
181+
182+
/// Returns the node \[u, p\[u\], .., lca(u, v), .., p\[v\], v\]\[k\], or None
183+
///
184+
/// # Complexity
185+
/// - Time: O(log n)
186+
/// - Space: O(1)
187+
pub fn kth_on_path(&self, u: usize, v: usize, k: usize) -> Option<usize> {
188+
let mut dst_side = [0; 2];
189+
self.path(u, v, |range, u_anc| dst_side[u_anc as usize] += range.len());
190+
if k < dst_side[0] {
191+
return self.kth_par(u, k);
192+
}
193+
let dst = dst_side[0] + dst_side[1] - !self.vals_edges as usize;
194+
if k <= dst {
195+
self.kth_par(v, dst - k)
196+
} else {
197+
None
198+
}
199+
}
129200
}

0 commit comments

Comments
 (0)