Initial commit

1 files changed, 153 insertions, 0 deletions

diff --git a/kattis-open-almostunionfind/src/main.rs b/kattis-open-almostunionfind/src/main.rs
new file mode 100644
index 0000000..6563043
--- /dev/null
+++ b/kattis-open-almostunionfind/src/main.rs
@@ -0,0 +1,153 @@
+// https://open.kattis.com/problems/almostunionfind
+
+use std::{
+    io::{stdin, Read},
+    iter::once,
+};
+
+enum Node {
+    Root { count: u32, sum: u64 },
+    Child(usize),
+    Invalid,
+}
+
+impl Node {
+    fn count(&self) -> u32 {
+        match self {
+            &Root { count, .. } => count,
+            _ => unreachable!(),
+        }
+    }
+}
+
+use Node::*;
+
+fn root(uf: &mut [Node], mut node: usize) -> usize {
+    let mut root = node;
+    while let Child(parent) = uf[root] {
+        root = parent
+    }
+    while let Child(parent) = uf[node] {
+        uf[node] = Child(root);
+        node = parent;
+    }
+    root
+}
+
+fn join(uf: &mut [Node], a: usize, b: usize) {
+    let mut a = root(uf, a);
+    let mut b = root(uf, b);
+
+    if a == b {
+        return;
+    }
+
+    if uf[a].count() < uf[b].count() {
+        let tmp = a;
+        a = b;
+        b = tmp;
+    }
+
+    let (b_count, b_sum) = match uf[b] {
+        Root { count, sum } => (count, sum),
+        _ => unreachable!(),
+    };
+    match &mut uf[a] {
+        Root { count, sum } => {
+            *count += b_count;
+            *sum += b_sum;
+        }
+        _ => unreachable!(),
+    }
+    uf[b] = Child(a);
+}
+
+fn mov(
+    uf: &mut [Node],
+    free_index: &mut usize,
+    current_index_of: &mut [usize],
+    a: usize,
+    a_val: u64,
+    b: usize,
+) {
+    let a_root = root(uf, a);
+    let b_root = root(uf, b);
+    if a_root == b_root {
+        return;
+    }
+    match &mut uf[a_root] {
+        Root { count, sum } => {
+            *count -= 1;
+            *sum -= a_val;
+        }
+        _ => unreachable!(),
+    }
+    match &mut uf[b_root] {
+        Root { count, sum } => {
+            *count += 1;
+            *sum += a_val;
+        }
+        _ => unreachable!(),
+    }
+    uf[*free_index] = Child(b_root);
+    current_index_of[a_val as usize] = *free_index;
+    *free_index += 1;
+}
+
+fn main() {
+    let mut input = String::new();
+    stdin().lock().read_to_string(&mut input).unwrap();
+    let mut ints = input
+        .trim()
+        .split(|c| c == ' ' || c == '\n')
+        .map(|i| i.parse::<usize>().unwrap());
+    let mut get = || ints.next();
+
+    while let Some(n) = get() {
+        let m = get().unwrap();
+
+        let mut uf: Vec<Node> = (0..n + m)
+            .map(|i| {
+                if i < n {
+                    Root {
+                        count: 1,
+                        sum: i as u64 + 1,
+                    }
+                } else {
+                    Invalid
+                }
+            })
+            .collect();
+        let mut free_index = n;
+        let mut current_index_of: Vec<usize> = (once(0).chain(0..n)).collect();
+
+        for _ in 0..m {
+            match get().unwrap() {
+                1 => {
+                    let (p, q) = (get().unwrap(), get().unwrap());
+                    join(&mut uf, current_index_of[p], current_index_of[q]);
+                }
+                2 => {
+                    let (p, q) = (get().unwrap(), get().unwrap());
+                    let (pi, qi) = (current_index_of[p], current_index_of[q]);
+                    mov(
+                        &mut uf,
+                        &mut free_index,
+                        &mut current_index_of,
+                        pi,
+                        p as u64,
+                        qi,
+                    );
+                }
+                _ => {
+                    let p = get().unwrap();
+                    let r = root(&mut uf, current_index_of[p]);
+                    match uf[r] {
+                        Root { count, sum } => println!("{} {}", count, sum),
+                        _ => unreachable!(),
+                    }
+                }
+            }
+        }
+    }
+}