Skip to main content

union_find/
quick_find.rs

1// Copyright 2016 union-find-rs Developers
2//
3// Licensed under the Apache License, Version 2.0, <LICENSE-APACHE or
4// http://apache.org/licenses/LICENSE-2.0> or the MIT license <LICENSE-MIT or
5// http://opensource.org/licenses/MIT>, at your option. This file may not be
6// copied, modified, or distributed except according to those terms.
7
8use crate::{Union, UnionFind, UnionResult};
9use std::iter::FromIterator;
10
11#[derive(Copy, Clone, Debug)]
12struct Payload<V> {
13    data: V,
14    link_last_child: usize,
15}
16
17/// Union-Find implementation with quick find operation.
18#[derive(Debug)]
19pub struct QuickFindUf<V> {
20    link_root: Vec<usize>,
21    link_sibling: Vec<usize>,
22    payload: Vec<Option<Payload<V>>>,
23}
24
25impl<V> Clone for QuickFindUf<V>
26where
27    V: Clone + Union,
28{
29    #[inline]
30    fn clone(&self) -> QuickFindUf<V> {
31        QuickFindUf {
32            link_root: self.link_root.clone(),
33            link_sibling: self.link_sibling.clone(),
34            payload: self.payload.clone(),
35        }
36    }
37
38    #[inline]
39    fn clone_from(&mut self, other: &QuickFindUf<V>) {
40        self.link_root.clone_from(&other.link_root);
41        self.link_sibling.clone_from(&other.link_sibling);
42        self.payload.clone_from(&other.payload);
43    }
44}
45
46impl<V: Union> UnionFind<V> for QuickFindUf<V> {
47    #[inline]
48    fn size(&self) -> usize {
49        self.payload.len()
50    }
51
52    #[inline]
53    fn insert(&mut self, data: V) -> usize {
54        let key = self.payload.len();
55        self.link_root.push(key);
56        self.link_sibling.push(key);
57        self.payload.push(Some(Payload {
58            data,
59            link_last_child: key,
60        }));
61        key
62    }
63
64    #[inline]
65    fn union(&mut self, key0: usize, key1: usize) -> bool {
66        let k0 = self.find(key0);
67        let k1 = self.find(key1);
68        if k0 == k1 {
69            return false;
70        }
71
72        // Temporary replace with dummy to move out the elements of the vector.
73        let Payload {
74            data: d0,
75            link_last_child: c0,
76        } = self.payload[k0].take().unwrap();
77        let Payload {
78            data: d1,
79            link_last_child: c1,
80        } = self.payload[k1].take().unwrap();
81
82        let (root, child_root, val, last) = match Union::union(d0, d1) {
83            UnionResult::Left(val) => (k0, k1, val, c0),
84            UnionResult::Right(val) => (k1, k0, val, c1),
85        };
86
87        self.link_sibling[last] = child_root;
88
89        let mut elem = child_root;
90        while self.link_sibling[elem] != elem {
91            debug_assert_eq!(self.link_root[elem], child_root);
92            self.link_root[elem] = root;
93            elem = self.link_sibling[elem];
94        }
95        debug_assert_eq!(self.link_root[elem], child_root);
96        self.link_root[elem] = root;
97
98        self.payload[root] = Some(Payload {
99            data: val,
100            link_last_child: elem,
101        });
102
103        true
104    }
105
106    #[inline]
107    fn find(&mut self, key: usize) -> usize {
108        self.link_root[key]
109    }
110
111    #[inline]
112    fn get(&mut self, key: usize) -> &V {
113        let root_key = self.find(key);
114        &self.payload[root_key].as_ref().unwrap().data
115    }
116
117    #[inline]
118    fn get_mut(&mut self, key: usize) -> &mut V {
119        let root_key = self.find(key);
120        &mut self.payload[root_key].as_mut().unwrap().data
121    }
122}
123
124impl<A: Union> FromIterator<A> for QuickFindUf<A> {
125    #[inline]
126    fn from_iter<T: IntoIterator<Item = A>>(iterator: T) -> QuickFindUf<A> {
127        let mut uf = QuickFindUf {
128            link_root: vec![],
129            link_sibling: vec![],
130            payload: vec![],
131        };
132        uf.extend(iterator);
133        uf
134    }
135}
136
137impl<A> Extend<A> for QuickFindUf<A> {
138    #[inline]
139    fn extend<T>(&mut self, iterable: T)
140    where
141        T: IntoIterator<Item = A>,
142    {
143        let len = self.payload.len();
144        let payload = iterable
145            .into_iter()
146            .zip(len..)
147            .map(|(data, link)| Payload {
148                data,
149                link_last_child: link,
150            })
151            .map(Some);
152        self.payload.extend(payload);
153
154        let new_len = self.payload.len();
155        self.link_root.extend(len..new_len);
156        self.link_sibling.extend(len..new_len);
157    }
158}