如何在Rust中按照预定义索引对Vec执行原地重排操作？

我们需要对大内存占用的Vec，按照给定的索引序列执行原地重排，避免额外内存开销，示例预期行为如下：

let i = vec![0, 3, 2, 1];
let mut v = vec!["a", "b", "c", "d"];
v.sort_by_indices(&i);

assert_eq!(v, &["a", "d", "c", "b"]);

实现思路 #

基于循环置换算法实现，仅需常数级额外内存，不需要分配与原Vec同规模的临时空间，完全满足大内存场景的要求。

版本1：不修改输入索引序列 #

需要额外的标记数组记录已处理位置，额外内存开销为原Vec长度的1/8（bool类型占1字节）：

use std::mem;

trait SortByIndices<T> {
    /// 按给定的索引序列原地重排Vec
    /// 
    /// # 注意
    /// 输入的indices必须是长度与self相等的合法排列：所有元素范围为`0..self.len()`且无重复
    fn sort_by_indices(&mut self, indices: &[usize]);
}

impl<T> SortByIndices<T> for Vec<T> {
    fn sort_by_indices(&mut self, indices: &[usize]) {
        assert_eq!(self.len(), indices.len(), "索引序列长度与Vec长度不匹配");
        let len = self.len();
        let mut visited = vec![false; len];

        for idx in 0..len {
            if visited[idx] || indices[idx] == idx {
                continue;
            }
            // 遍历当前循环置换链
            let mut current_idx = idx;
            let mut current_val = mem::take(&mut self[current_idx]);
            loop {
                let target_idx = indices[current_idx];
                if visited[target_idx] {
                    break;
                }
                // 交换当前值与目标位置值
                mem::swap(&mut current_val, &mut self[target_idx]);
                visited[current_idx] = true;
                current_idx = target_idx;
            }
            // 把循环末尾值放回起始位置
            self[idx] = current_val;
            visited[idx] = true;
        }
    }
}

版本2：零额外堆内存（允许修改输入索引序列） #

如果场景允许修改输入的索引数组，可以直接在索引数组上做已处理标记，完全消除额外堆内存开销：

use std::mem;

trait SortByIndices<T> {
    /// 按给定的索引序列原地重排Vec，会修改输入的索引序列做处理标记
    /// 
    /// # 注意
    /// 输入的indices必须是长度与self相等的合法排列：所有元素范围为`0..self.len()`且无重复
    fn sort_by_indices(&mut self, indices: &mut [usize]);
}

impl<T> SortByIndices<T> for Vec<T> {
    fn sort_by_indices(&mut self, indices: &mut [usize]) {
        assert_eq!(self.len(), indices.len(), "索引序列长度与Vec长度不匹配");
        let len = self.len();

        for idx in 0..len {
            if indices[idx] == usize::MAX || indices[idx] == idx {
                continue;
            }
            let mut current_idx = idx;
            let mut current_val = mem::take(&mut self[current_idx]);
            loop {
                let target_idx = indices[current_idx];
                if indices[target_idx] == usize::MAX {
                    break;
                }
                mem::swap(&mut current_val, &mut self[target_idx]);
                indices[current_idx] = usize::MAX;
                current_idx = target_idx;
            }
            self[idx] = current_val;
            indices[idx] = usize::MAX;
        }
    }
}

特性说明 #

• 兼容所有Rust合法类型，不需要元素实现Clone、Copy trait
• 全程无同规模临时内存分配，版本2仅占用栈上常数级内存，完全适配超大内存Vec场景
• 若不需要极致性能，可移除代码中的unsafe块（示例代码已默认使用安全索引访问），Rust自动边界检查的性能损耗可忽略

内容的提问来源于stack exchange，提问作者Michael Hall