6
\$\begingroup\$

Related: Find all distinct subsets that sum to a given number

This code is supposed to efficiently generate all subsets of a list such that the subset's values sum to a particular target value. For example, given the input list [5, 3, 2, 1] and the target 8, it should produce [5, 3] and [5, 2, 1]. Given the same input list and the target 6, it should produce [5, 1] and [3, 2, 1].

I couldn't locate any academic research on the topic, so I invented this algorithm. I'm pretty confident that it's efficient, but would be very interested in potential improvements. I'd also love any suggestions that could reduce use of unsafe or remove the necessity to pin the struct.

Rust playground (with tests)

pub trait Permutable: Copy + Ord + Sub<Output = Self> {}
impl<T: Copy + Ord + Sub<Output = Self>> Permutable for T {}

type Solution<Subset> = Vec<Option<Subset>>;

/// A `BoundedPermutationGenerator` efficiently generates distinct subsets of an input list of items,
/// where each subset sums to a particular value.
///
/// Its output is `Solution`: a `Vec<Option<Subset>>`. A non-`None` entry in the `Solution`
/// indicates that the corresponding value in `items` should be assigned to that subset.
///
/// It can accept a partial solution with [`from_solution`][Self::from_solution]. In that case,
/// [`next_solution_for`][Self::next_solution_for] will never select a value already used for a
/// different subset.
///
/// Note that the constructors always produce `Pin<Box<Self>>`. This is an implementation detail.
/// To function efficiently, this contains (many) references to an area of memory used as a scratch
/// space. If it were possible to use an external slice for this, then we wouldn't need to pin.
/// However, we strongly desire that it's possible to construct a free-standing generator from this;
/// that implies that this struct needs to own its working memory. That means that those
/// references must be internal, which in turn means that we need to ensure that this struct never
/// moves so that the internal references remain valid.
#[derive(Debug, Clone)]
pub struct BoundedPermutationGenerator<'a, T, Subset> {
    _lifetime: PhantomData<&'a ()>,
    scratch_space: Solution<Subset>,
    inner: Inner<T, Subset>,
}

impl<'a, T, Subset> BoundedPermutationGenerator<'a, T, Subset>
where
    T: Permutable,
    Subset: Copy + Eq,
{
    /// Create a new `BoundedPermutationGenerator` from a list of items.
    ///
    /// # Preconditions
    ///
    /// - `packges` must be reverse-sorted.
    pub fn new(items: &'a [T], target_sum: T) -> Result<Pin<Box<Self>>, Error> {
        Self::from_solution(items, target_sum, vec![None; items.len()])
    }

    /// Create a new `BoundedPermutationGenerator` from a list of items and an existing partial solution.
    ///
    /// # Preconditions
    ///
    /// - `items` must be reverse-sorted.
    /// - `solution.len()` must equal `items.len()`.
    pub fn from_solution(
        items: &'a [T],
        target_sum: T,
        solution: Solution<Subset>,
    ) -> Result<Pin<Box<Self>>, Error> {
        if solution.len() != items.len() {
            return Err(Error::WrongSolutionSize(solution.len(), items.len()));
        }
        if !items.windows(2).all(|window| window[1] <= window[0]) {
            return Err(Error::ItemsNotSorted);
        }
        let mut temporary = Vec::new();
        // we'll update the `scratch_space` field to be meaningful once the outer is pinned
        let bpg = BoundedPermutationGenerator {
            _lifetime: PhantomData,
            scratch_space: solution,
            inner: Inner {
                items: items as _,
                scratch_space: temporary.as_mut_slice() as _,
                target_sum,
                idx: 0,
                child: None,
            },
        };
        let mut boxed = Box::pin(bpg);
        let scratch_space = boxed.scratch_space.as_slice() as *const _ as _;

        // safe because modifying a field never moves the whole struct
        unsafe {
            let mut_ref: Pin<&mut Self> = Pin::as_mut(&mut boxed);
            Pin::get_unchecked_mut(mut_ref).inner.scratch_space = scratch_space;
        }

        Ok(boxed)
    }

    /// Generate the next valid layout for members of this subset.
    ///
    /// Each solution requires an allocation and data-copying proportional to `self.scratch_space`.
    pub fn next_solution_for(&mut self, subset: Subset) -> Option<Solution<Subset>> {
        self.inner.next_solution_for(subset)
    }

    /// Transform into an iterator over the remaining solutions of this generator.
    pub fn into_iter(self: Pin<Box<Self>>, subset: Subset) -> Iter<'a, T, Subset> {
        Iter { bpg: self, subset }
    }
}

/// The inner structure contains all the actual implementation details of the solution generator.
///
/// It's a separate, private struct because it uses raw pointers instead of normal references.
/// This is because it's recursive, and rustc can't figure out an appropriate lifetime otherwise.
///
/// Imagine that we hadn't separated the lifetime from the references: it would look like
///
/// ```
/// pub struct BoundedPermutationGenerator<'a, T, Subset> {
///   items: &'a [T],
///   scratch_space: &'a mut [Option<Subset>],
///   item_idx: usize,
///   target_sum: T,
///   child: Option<Box<BoundedPermutationGenerator<'a, T, Subset>>>,
/// }
/// ```
///
/// The problem is the `child` field: because we've defined it to be `'a`, then any borrow _must_
/// last for that long, which doesn't work with the recursive strategy that we want to use. However,
/// we don't have access to any other lifetime which we can use.
///
/// Splitting the lifetime away means that we have to do a little more work to ensure that everything
/// stays safe, but it also means that this minimal-copy approach is possible at all.
#[derive(Debug, Clone)]
struct Inner<T, Subset> {
    /// a reverse-sorted list of available items.
    items: *const [T],
    /// working space containing the current state of the partial solution.
    scratch_space: *mut [Option<Subset>],
    /// an index into `self.items` and `self.scratch_space`.
    idx: usize,
    /// keeps track of the value we're looking for at this depth of recursion.
    target_sum: T,
    /// if present, a recursive child assuming that the current item is part of the solution.
    child: Option<Box<Inner<T, Subset>>>,
}

impl<T, Subset> Inner<T, Subset>
where
    T: Permutable,
    Subset: Copy + Eq,
{
    /// Private access to `self.items` as a slice.
    ///
    /// Safe because the only way to construct a `BoundedPermutationGenerator` requires a valid slice,
    /// and we never edit the pointer.
    fn items(&self) -> &[T] {
        unsafe { &*self.items }
    }

    /// Private access to `self.scratch_space` as a slice.
    ///
    /// Safe because the only way to construct a `BoundedPermutationGenerator` requires a valid slice,
    /// and we never edit the pointer.
    fn scratch_space(&self) -> &[Option<Subset>] {
        unsafe { &*self.scratch_space }
    }

    /// Private mutable access to `self.scratch_space` as a slice.
    ///
    /// Safe because the only way to construct a `BoundedPermutationGenerator` requires a valid slice,
    /// and we never edit the pointer.
    fn scratch_space_mut(&self) -> &mut [Option<Subset>] {
        unsafe { &mut *self.scratch_space }
    }

    /// Create a child generator which can be used to recursively seek solutions.
    fn child(&mut self) -> Box<Self> {
        Box::new(Self {
            items: self.items,
            scratch_space: self.scratch_space,
            target_sum: self.target_sum - self.items()[self.idx],
            idx: self.idx + 1,
            child: None,
        })
    }

    /// Recursively generate the next valid layout for members of this subset.
    ///
    /// Each solution requires an allocation and data-copying proportional to `self.scratch_space`.
    ///
    /// # Method of operation
    ///
    /// - if we have a child, produce its next solution
    /// - if we don't have a child, for each index available to us:
    ///   - if it's greater than the target, then just try the next one
    ///   - if it's equal to the target, then we know the solution is complete
    ///   - if it's less than the target, create a child which will attempt to generate solutions
    ///     assuming that the current item is a member of the solution set
    ///
    /// Because this is recursive and cleans up after itself, the stack provides efficient backtracking.
    fn next_solution_for(&mut self, subset: Subset) -> Option<Solution<Subset>> {
        let mut solution = None;
        while solution.is_none() && self.idx < self.items().len() {
            match self.child {
                None => {
                    if let Some(existing_subset) = self.scratch_space()[self.idx] {
                        if existing_subset == subset {
                            // we've re-entered after returning a valid solution.
                            // To avoid infinite loops, unset this value and try the next.
                            self.scratch_space_mut()[self.idx] = None;
                        }
                        // otherwise never overwrite a previously-set member of the subset layout.
                        // this property is essential for composability.
                        self.idx += 1;
                        continue;
                    }

                    match self.items()[self.idx].cmp(&self.target_sum) {
                        Ordering::Greater => {
                            // no luck; try the next one
                            self.idx += 1;
                        }
                        Ordering::Equal => {
                            // we've identified a valid solution.
                            self.scratch_space_mut()[self.idx] = Some(subset);
                            solution = Some(self.scratch_space().to_vec());
                        }
                        Ordering::Less => {
                            // recursively try different subsets assuming this item is a member of
                            // the solution.
                            self.scratch_space_mut()[self.idx] = Some(subset);
                            self.child = Some(self.child());
                        }
                    }
                }
                Some(ref mut child) => {
                    match child.next_solution_for(subset) {
                        Some(inner_solution) => {
                            // while the child produces solutions, just pass them along.
                            solution = Some(inner_solution);
                        }
                        // If the child stops producing values, unset it; the next iteration of the
                        // main loop will clean up the rest.
                        None => self.child = None,
                    }
                }
            };
        }
        solution
    }
}

#[derive(Debug, PartialEq, Eq, thiserror::Error)]
pub enum Error {
    #[error("`items` input was not sorted")]
    ItemsNotSorted,
    #[error("Solution has {0} entries but must have {1} to match items")]
    WrongSolutionSize(usize, usize),
}
\$\endgroup\$
2

0

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service and acknowledge you have read our privacy policy.

Browse other questions tagged or ask your own question.