1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
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!(),
}
}
}
}
}
}
|
