1
use std::fmt;
2

            
3
use crate::Partition;
4

            
5
/// Defines a partition based on an explicit indexing of elements to their block
6
/// number.
7
#[derive(Debug)]
8
pub struct IndexedPartition {
9
    partition: Vec<usize>,
10

            
11
    num_of_blocks: usize,
12
}
13

            
14
impl IndexedPartition {
15
    /// Create a new partition where all elements are in a single block.
16
397
    pub fn new(num_of_elements: usize) -> IndexedPartition {
17
397
        IndexedPartition {
18
397
            partition: vec![0; num_of_elements],
19
397
            num_of_blocks: 1,
20
397
        }
21
397
    }
22

            
23
    /// Create a new partition with the given partitioning.
24
29
    pub fn with_partition(partition: Vec<usize>, num_of_blocks: usize) -> IndexedPartition {
25
29
        IndexedPartition {
26
29
            partition,
27
29
            num_of_blocks,
28
29
        }
29
29
    }
30

            
31
    /// Iterates over the elements in the partition.
32
56
    pub fn iter(&self) -> impl Iterator<Item = usize> + '_ {
33
56
        self.partition.iter().copied()
34
56
    }
35

            
36
    /// Sets the block number of the given element
37
5067577
    pub fn set_block(&mut self, element_index: usize, block_number: usize) {
38
5067577
        // TODO: This assumes that the blocks are dense, otherwise it overestimates the number of blocks.
39
5067577
        self.num_of_blocks = self.num_of_blocks.max(block_number + 1);
40
5067577

            
41
5067577
        self.partition[element_index] = block_number;
42
5067577
    }
43
}
44

            
45
/// Combines two partitions into a new partition.
46
56
pub fn combine_partition(left: IndexedPartition, right: &impl Partition) -> IndexedPartition {
47
56
    let mut combined_partition = IndexedPartition::new(left.partition.len());
48

            
49
282488
    for (element_index, block) in left.partition.iter().enumerate() {
50
282488
        let new_block = right.block_number(*block);
51
282488

            
52
282488
        combined_partition.set_block(element_index, new_block);
53
282488
    }
54

            
55
56
    combined_partition
56
56
}
57

            
58
/// Reorders the blocks of the given partition according to the given permutation.
59
56
pub fn reorder_partition<P>(partition: IndexedPartition, permutation: P) -> IndexedPartition
60
56
where
61
56
    P: Fn(usize) -> usize,
62
56
{
63
56
    let mut new_partition = IndexedPartition::new(partition.len());
64

            
65
282488
    for (element_index, block) in partition.iter().enumerate() {
66
282488
        new_partition.set_block(element_index, permutation(block));
67
282488
    }
68

            
69
56
    new_partition
70
56
}
71

            
72
impl PartialEq for IndexedPartition {
73
54
    fn eq(&self, other: &Self) -> bool {
74
54
        self.equal(other)
75
54
    }
76
}
77

            
78
impl fmt::Display for IndexedPartition {
79
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
80
        write!(f, "{{ ")?;
81

            
82
        let mut first = true;
83

            
84
        for block_index in 0..self.partition.len() {
85
            // Print all elements with the same block number.
86
            let mut first_block = true;
87
            for (element_index, _) in self.iter().enumerate().filter(|(_, value)| *value == block_index) {
88
                if !first_block {
89
                    write!(f, ", ")?;
90
                } else {
91
                    if !first {
92
                        write!(f, ", ")?;
93
                    }
94

            
95
                    write!(f, "{{")?;
96
                }
97

            
98
                write!(f, "{}", element_index)?;
99
                first_block = false;
100
            }
101

            
102
            if !first_block {
103
                write!(f, "}}")?;
104
                first = false;
105
            }
106
        }
107

            
108
        write!(f, " }}")
109
    }
110
}
111

            
112
impl Partition for IndexedPartition {
113
11884654092
    fn block_number(&self, state_index: usize) -> usize {
114
11884654092
        self.partition[state_index]
115
11884654092
    }
116

            
117
981382
    fn num_of_blocks(&self) -> usize {
118
981382
        self.num_of_blocks
119
981382
    }
120

            
121
316436
    fn len(&self) -> usize {
122
316436
        self.partition.len()
123
316436
    }
124
}