Iterating over all possible partitions of a given set

user_4322240 0
Viewed 241

Given a set represented by a vector, e.g. [1,2,3], I want to iterate over all possible partitions (the order does not matter to me):

[[1], [2], [3]]
[[1,2], [3]]
[[1,3], [2]]
[[1], [2,3]]
[[1,2,3]]

I did not find a way to do this in the itertools package, neither did I find anything elsewhere. The closest I have seen was this post on StackOverflow, however in that post not all partitions are generated, e.g. [[1,3],[2]] in this case. Is there maybe some existing package to do this? Or how could I achieve that?

Edit: To clarify:

  • With "set" and "partition" I mean the mathematical notions, i.e.: a set X is a partition of a set A iff every element of X is a non-empty subset of A and every element of A appears in exactly one element of X.
  • With "representation of a set by a vector" I mean a vector that contains each element of the set exactly once, in any order. In particular, it won't contain duplicate elements.
  • What I actually want is an efficient iterator, which yields (a representation of) each possible partition exactly once. I.e. I am interested in some method create_iterator(v: Vec<T>) -> impl Iterator<Item = Vec<Vec<T>>>. I would prefer not to generate a vector of all partitions, since this has the drawbacks that it needs more space and, moreover, using the iterator I can break, which can drastically reduce runtime.
  • An example use-case could look like this:
for x in create_iterator(vec![1,2,3]) {
    println!("{:?}", x);
}
rust algorithm
2 Answers

This may not be the most efficient in terms of heap allocation, but it works! Here's the playground.


fn get_partitions(input: &mut Vec<u64>) -> Vec<Vec<Vec<u64>>> {
    if input.len() == 0 {
        return Vec::new();
    }
    if input.len() == 1 {
        return vec![vec![vec![input[0]]]];
    }
    else {
        let a = input.pop().unwrap();
        let partitions = get_partitions(input);
        let mut output = Vec::new();
        for part in partitions {
            let mut tmp = part.to_vec();
            tmp.push(vec![a]);
            output.push(tmp);
            for idx in 0..part.len() {
                let mut tmp = part.to_vec();
                tmp[idx].push(a);
                output.push(tmp);
            }
        }
        return output;
    }
}

fn main() {
    let mut input = vec![1,2,3];
    let output = get_partitions(&mut input);
    println!("{:?}", output);
}

I found the answer here, which outlines the algorithm.

user_126580930

In the end I just implemented the highest-ranked answer here. I am still new to the language, and hence the implementation may not be the best. In particular, the lifetime parameters were a bit messy. So, comments are highly appreciated! :)

use itertools::zip;

// Based on https://stackoverflow.com/a/30898130/4322240. Read first if something here is unclear!

// Representation of a state in which the iterator can be.
// elements: contains the elements to be partitioned.
// index_map: in terms of the linked post, index_map[i] = p_i (but 0-indexed)
// initialized: marks if the Partition element was just initialized and hence
//     was not generated by next()
pub struct Partition<'a, T> where T: Copy {
    elements: &'a Vec<T>,
    index_map: Vec<usize>,
    initialized: bool, 
}
impl<'a, T> Partition<'a, T> where T: Copy {
    fn smallest_viable(v: &'a Vec<T>) -> Partition<'a, T> {
        Partition::<T> { elements: v, index_map: vec![0; v.len()], initialized: true }
    }
    fn is_viable(&self) -> bool {
        self.elements.len() == self.index_map.len()
        && (self.index_map.is_empty() || self.index_map[0] == 0)
        && (1..self.index_map.len())
                .all(|x| (0..x).any(|y| self.index_map[x] <= self.index_map[y] + 1))
    }
    fn is_incrementable(&self, index: usize) -> bool {
        index <= self.index_map.len()
        && index > 0
        // p_i <= n does not need to be checked, as it is guaranteed if the invariant
        // "for all j > 0 there exists i < j s.t. p_j <= p_i + 1" is satisfied
        && (0..index).any(|x| self.index_map[x] == self.index_map[index])
    }
    fn increment(&mut self, index: usize) {
        self.index_map[index] += 1;
        for x in index + 1..self.index_map.len() {
            self.index_map[x] = 0;
        }
    }
    fn create_partition(&self) -> Vec<Vec<T>> {
        if let Some(&max_index) = self.index_map.iter().max(){
            return (0..=max_index).map(|x| zip(self.elements, &self.index_map)
                    .filter(|(_e,i)| **i == x).map(|(e,_i)| *e).collect())
                    .collect();
        } else {
            return Vec::new();
        }
    }
}
impl<'a, T> Iterator for Partition<'a, T> where T: Copy {
    type Item = Vec<Vec<T>>;

    fn next(&mut self) -> Option<Self::Item> {
        if self.initialized {
            self.initialized = false;
            return Some(self.create_partition());
        }
        for index in (0..self.index_map.len()).rev() {
            if self.is_incrementable(index) {
                self.increment(index);
                return Some(self.create_partition());
            }
        }
        return None;
    }
}

pub fn partitions<'a, T>(v: &'a Vec<T>) -> Partition<'a, T> where T: Copy {
    Partition::smallest_viable(v)
}

Example usage:

fn main() {
    for x in partitions(&vec![7,8,9]) {
        println!("{:?}", x);
    }
}

outputs

[[7, 8, 9]]
[[7, 8], [9]]
[[7, 9], [8]]
[[7], [8, 9]]
[[7], [8], [9]]
user_43222400
Related