有没有一种方法可以根据谓词来排除向量的部分元素？

Question

我试图根据谓词从向量中删除一些元素，并收集结果。以下是一个（不起作用的）示例和预期结果：

let mut v: Vec<i32> = vec![1, 2, 3, 4, 5, 6];

let drained: Vec<i32> = v.iter().filter(|e| (*e) % 2 == 0).drain(..).collect();

assert_eq!(v, vec![1, 3, 5]);
assert_eq!(drained, vec![2, 4, 6]);

这会导致错误。

error[E0599]: no method named `drain` found for type `std::iter::Filter<std::slice::Iter<'_, i32>, [closure@src/main.rs:4:45: 4:62]>` in the current scope
 --> src/main.rs:4:64
  |
4 |     let drained: Vec<i32> = v.iter().filter(|e| (*e) % 2 == 0).drain(..).collect();
  |                                                                ^^^^^

我看了几个替代方案，但似乎都不能满足我的需求：

• Vec::retain 移除向量中的元素，但不返回已删除元素的所有权。

• v.drain(..).filter(condition).collect() 返回正确的 drained 值，但会清空整个向量。

- WorldSEnder

---

你尝试过使用 to_vec 将数组转换吗？ let drained: Vec<i32> = v.iter().filter(|e| (*e) % 2 == 0).to_vec().drain(..).collect(); - aug2uag

Answer 1

在稳定的 Rust 1.33.0 中没有这个功能。不过有一个名为 drain_filter 的不稳定 nightly 特性可以做到你想要的：

#![feature(drain_filter)]

fn main() {
    let mut v: Vec<i32> = vec![1, 2, 3, 4, 5, 6];

    let drained: Vec<i32> = v.drain_filter(|&mut e| e % 2 == 0).collect();

    assert_eq!(v, vec![1, 3, 5]);
    assert_eq!(drained, vec![2, 4, 6]);
}

作为一个稳定的解决办法，你可以使用Iterator::partition，但它不会重用内存：

fn main() {
    let v: Vec<i32> = vec![1, 2, 3, 4, 5, 6];

    let (drained, v): (Vec<_>, Vec<_>) = v.into_iter().partition(|&e| e % 2 == 0);

    assert_eq!(v, vec![1, 3, 5]);
    assert_eq!(drained, vec![2, 4, 6]);
}

- Shepmaster

---

供以后参考：RFC和跟踪问题。 - dj_bushido

Answer 2

我可以提供几种其他方法来做这件事，还有我的基准测试结果。

注意： 我会根据以下几个标准比较所有方法：

1. 是否支持外部真实源（我的使用情况）。例如，Vec::retain支持此功能，这意味着您可以编写以下代码：

// conditions: &[bool]
assert_eq!(conditions.len(), my_vec.len());
let cond = conditions.iter().copied();
my_vec.retain(move|_|cond.next().unwrap());

2. 这个方法是否被第三方的 Vec 支持，比如 ArrayVec、TinyVec、SmallVec、FixedSliceVec 等。
3. 它的速度快吗。

那么，让我们开始吧。

先排序再分割切片

特点：

1. 支持外部真相来源 — 不支持 ❌。会以未指定的顺序调用闭包 O(n log n) 次，因此只支持仅从值直接计算出的谓词。
2. 第三方支持 — 非常好 ✅。您可以将其用于任何可转换为可变切片的内容。
3. 速度快吗 — 在通用情况下，不快 ❌。它的运行时间为 O(n log n)，而其他方法的运行时间为 O(n)。但是，如果对于您来说保留原始相对顺序并不重要，您可以使用 sort_unstable_by_key，它根本不会分配内存，这在某些情况下可能是最快的。

实现：

v.sort_by_key(|x| predicate(x));
let split_pos = v.partition_point(|x| predicate(x));
let (false_slice, true_slice) = v.split_at_mut(split_pos)

Vec::drain_filter

1. 支持外部真相源 — 是 ✅。按原始顺序仅访问项目一次。
2. 第三方支持 — 不存在 ❌。此外，您甚至不能在稳定的 Rust 中使用它，它的跟踪问题 已经饱受五年的争论之苦（截至2022年7月）。
3. 速度快吗 — 是 ✅。

代码

let removed_items: Vec<_> = v.drain_filter(|x| predicate(x)).collect();

MaybeUninit技巧使用unsafe代码

嗯，我自己写的。

特点：

1. 支持外部真相源 - 是 ✅。按原始顺序仅访问项目一次。
2. 第三方支持 - 支持 ✅。请注意，您必须审核他们对retain的实现，以确保它本身不会出现问题。
3. 速度快 - 是 ✅。在我的基准测试中，它比Vec::drain_filter更快。

此实现有两个假设：

1. Vec<T>和Vec<MaybeUninit<T>>的内部布局是相同的。这是没有理由的，因为T和MaybeUninit<T>的内存布局是相同的，但我也使用了断言来验证这一点。请注意，断言将被编译器优化器移除，因为它们总是为真。
2. retain本身不会引发panic（它只能从元素删除或谓词中传播panic）。这对于std::vec::Vec是正确的，但您需要确保第三方包也是如此。

算法很简单：

1. 将初始向量Vec<T>重新解释为Vec<MaybeUninit<T>>；
2. 将我们的谓词包装成新的谓词，将我们想要删除的项移动到外部存储中；
3. 让Vec::retain处理删除项目。

此外，使用MaybeUninit和unsafe的唯一原因是避免双重释放，因此如果您的元素实现了Copy，则可以在安全的Rust中实现此算法。 但是，在这种情况下，您可以只使用filter(...).collect()和retain，几乎具有相同的性能。

因此，以下是代码及其所有注释，说明它为什么是安全的（请注意，我没有使用sanitizer或Miri进行测试，因此使用时需自行承担风险）：

/// Returns removed values.
fn retain_unsafe_generic<T: Sized>(
    v: &mut Vec<T>,
    mut which_to_keep: impl FnMut(&T) -> bool,
) -> Vec<T> {
    use std::mem::{transmute, MaybeUninit};

    /// # Safety
    /// Caller must ensure that if it makes living copies of inner items,
    /// those items is removed from original vec before original reference became usable again.
    unsafe fn as_uninits<T: Sized>(v: &mut Vec<T>) -> &mut Vec<MaybeUninit<T>> {
        let orig_ptr = v.as_ptr();
        let orig_cap = v.capacity();
        let orig_size = v.len();

        let v: &mut Vec<MaybeUninit<T>> = unsafe {
            // Safety: since `MaybeUninit` has same memory layout
            // as wrapped type, we assume that we can treat vec with T
            // as MaybeUninit<T>. This assumption checked by asserts below.
            //
            // Lifetimes of elements must be correctly enforced by caller.
            transmute(v)
        };
        // Check if layout of Vec with different element type remains same.
        assert_eq!(v.len(), orig_size);
        assert_eq!(v.capacity(), orig_cap);
        assert_eq!(v.as_ptr(), orig_ptr.cast());

        v
    }

    let mut res: Vec<T> = Vec::with_capacity(v.len());
    let v = unsafe {
        // Safety: we keep result reference only in `retain` call.
        // We would remove all moved elements using retain.
        as_uninits(v)
    };
    v.retain(
        // Safety: `Vec::retain` would remove all items which values we moved into `res`.
        // It wouldn call `drop::<T>` for removed values
        // because `MaybeUninit` never drops wrapped values.
        |x| unsafe {
            // Safety: it is safe because `Vec::retain` visits elements sequentally
            // so we haven't moved value from `x` yet.
            // https://doc.rust-lang.org/std/vec/struct.Vec.html#method.retain
            let val = &*x.as_ptr();
            if which_to_keep(val) {
                return true;
            }
            res.reserve(1);
            // Any panic before this place is safe because
            // 1. We didn't moved value from `x` yet;
            // 2. In case of panic in predicate, `Vec::retain` preserve current value.

            // Here we could probably use `Vec::push`
            // but compiler currently fails to remove capacity check in `Vec::push`
            // so this function became slower than `Vec::drain_filter`
            // https://godbolt.org/z/7fhnnMh46
            // And `Vec::push(x.assume_init_read())` is unsafe operation too anyway.

            let old_len = res.len();
            // Safety: We just allocated memory for this place.
            let dst = res.as_mut_ptr().add(old_len);
            // Safety: since we cannot panic until the end of closure
            // and `Vec::retain` wouldn't panic and would remove `x`,
            // making bitwise copy of `x` is safe.
            x.as_ptr().copy_to_nonoverlapping(dst, 1);
            // Safety: we just wrote additional value.
            res.set_len(old_len + 1);

            false
        },
    );

    res
}

基准测试

基准测试的代码很长，所以这里提供到代码片段链接: https://gist.github.com/AngelicosPhosphoros/7ee482316bc1c83945f88308954e0d7e 它尝试使用我列出的所有三个算法从Vec中分离出奇数。

结果:

算法	混合	首先奇数	首先偶数
sort-split	465微秒	35微秒	10微秒
drain_filter	26微秒	24微秒	22.5微秒
retain-uninit	17微秒	21微秒	19微秒

正如您所看到的，在除了 sort-split 没有实际作用的情况下，retain 的使用在所有情况下都胜出。 这主要是因为 Vec::retain 在多年的优化中已经得到了严格的优化。

- Angelicos Phosphoros

Answer 3

文档说明Vec.retain将原地操作，并按顺序访问每个元素，仅一次。

fn drain_where<T, Pred : Fn(&T) -> bool>(source: &mut Vec<T>, pred: Pred) -> Vec<T>
  where T : Copy {
  let mut drained: Vec<T> = Vec::new();
  
  source.retain(|item| {
    if pred(item) { drained.push(*item); false } else { true }
  });
  
  drained
}

- Ivar Svendsen

---

它只适用于“Copy”类型，因此通常没有什么用处。 - nirvana-msu