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
154
155
156
157
158
159
160
161
162
163
164

use crate::arraymap::U16ArrayMap;
use std::cmp::{max, min};

const NO_CONNECTED_COMPONENT: u16 = 0;

#[derive(Clone, Debug, PartialEq, PartialOrd, Eq, Ord, Hash)]
pub struct ConnectedComponentGraph {
    // Mapping from nodes to their connected component id
    node_connected_component: U16ArrayMap,
    // Mapping from original connected component id to the one they've been merged with
    merged_connected_components: U16ArrayMap,
    // Size of each connected component in the graph
    connected_component_size: U16ArrayMap,
    num_connected_components: usize,
}

impl ConnectedComponentGraph {
    pub fn new(max_nodes: usize) -> ConnectedComponentGraph {
        let mut result = ConnectedComponentGraph {
            node_connected_component: U16ArrayMap::new(0, max_nodes),
            merged_connected_components: U16ArrayMap::new(
                (NO_CONNECTED_COMPONENT + 1) as usize,
                max_nodes,
            ),
            connected_component_size: U16ArrayMap::new(
                (NO_CONNECTED_COMPONENT + 1) as usize,
                max_nodes,
            ),
            num_connected_components: 0,
        };
        for i in result.merged_connected_components.keys() {
            result.merged_connected_components.insert(i, i as u16);
        }

        result
    }

    pub fn create_connected_component(&mut self) -> u16 {
        self.num_connected_components += 1;

        NO_CONNECTED_COMPONENT + self.num_connected_components as u16
    }

    pub fn add_node(&mut self, node: usize, connected_component: u16) {
        assert!(connected_component <= self.num_connected_components as u16);
        assert_eq!(
            self.node_connected_component.get(node),
            NO_CONNECTED_COMPONENT
        );
        let canonical = self.canonical_component_id(connected_component as usize);
        self.node_connected_component.insert(node, canonical as u16);
        self.connected_component_size.increment(canonical);
    }

    pub fn swap(&mut self, node1: usize, node2: usize) {
        self.node_connected_component.swap(node1, node2);
    }

    pub fn contains(&self, node: usize) -> bool {
        self.node_connected_component.get(node) != NO_CONNECTED_COMPONENT
    }

    pub fn remove_node(&mut self, node: usize) {
        let connected_component =
            self.canonical_component_id(self.node_connected_component.get(node) as usize);
        if connected_component == NO_CONNECTED_COMPONENT as usize {
            return;
        }
        self.connected_component_size
            .decrement(connected_component as usize);
        self.node_connected_component
            .insert(node, NO_CONNECTED_COMPONENT);
    }

    pub fn get_node_in_largest_connected_component(
        &self,
        start_node: usize,
        end_node: usize,
    ) -> usize {
        let mut max_size = 0;
        let mut largest_connected_component = NO_CONNECTED_COMPONENT as usize;
        for i in 1..=self.num_connected_components {
            let size = self.connected_component_size.get(i as usize);
            if size > max_size {
                max_size = size;
                largest_connected_component = i;
            }
        }
        assert_ne!(largest_connected_component, NO_CONNECTED_COMPONENT as usize);

        // Find a node (column) in that connected component
        (start_node..end_node)
            .find(|node| {
                self.canonical_component_id(self.node_connected_component.get(*node) as usize)
                    == largest_connected_component
            })
            .unwrap()
    }

    // This function will implicitly create any missing nodes
    pub fn add_edge(&mut self, node1: usize, node2: usize) {
        let connected_component1 =
            self.canonical_component_id(self.node_connected_component.get(node1) as usize);
        let connected_component2 =
            self.canonical_component_id(self.node_connected_component.get(node2) as usize);

        if connected_component1 == NO_CONNECTED_COMPONENT as usize
            && connected_component2 == NO_CONNECTED_COMPONENT as usize
        {
            // Create a new connected component
            let connected_component_id = self.create_connected_component();
            self.node_connected_component
                .insert(node1, connected_component_id);
            self.node_connected_component
                .insert(node2, connected_component_id);
            self.connected_component_size
                .insert(connected_component_id as usize, 2);
        } else if connected_component1 == NO_CONNECTED_COMPONENT as usize {
            self.connected_component_size
                .increment(connected_component2);
            self.node_connected_component
                .insert(node1, connected_component2 as u16);
        } else if connected_component2 == NO_CONNECTED_COMPONENT as usize {
            self.connected_component_size
                .increment(connected_component1);
            self.node_connected_component
                .insert(node2, connected_component1 as u16);
        } else if connected_component1 != connected_component2 {
            // Merge into the lowest to keep chains short
            let merge_to = min(connected_component1, connected_component2);
            let merge_from = max(connected_component1, connected_component2);
            let to_size = self.connected_component_size.get(merge_to as usize);
            let from_size = self.connected_component_size.get(merge_from as usize);
            self.connected_component_size.insert(merge_from as usize, 0);
            self.connected_component_size
                .insert(merge_to as usize, to_size + from_size);
            self.merged_connected_components
                .insert(merge_from, merge_to as u16);
        }
    }

    fn canonical_component_id(&self, mut id: usize) -> usize {
        if id == NO_CONNECTED_COMPONENT as usize {
            return id;
        }
        while self.merged_connected_components.get(id) as usize != id {
            id = self.merged_connected_components.get(id) as usize;
        }

        id
    }

    pub fn reset(&mut self) {
        for i in 1..=self.num_connected_components {
            self.connected_component_size.insert(i, 0);
            self.merged_connected_components.insert(i, i as u16);
        }
        self.num_connected_components = 0;
        for i in self.node_connected_component.keys() {
            self.node_connected_component
                .insert(i, NO_CONNECTED_COMPONENT);
        }
    }
}