Как избежать выделения памяти в Iterator :: flat_map?

У меня есть Vec целых чисел, и я хочу создать новый Vec, содержащий эти целые числа и квадраты этих целых чисел. Я мог бы сделать это в обязательном порядке:

let v = vec![1, 2, 3];
let mut new_v = Vec::new(); // new instead of with_capacity for simplicity sake.
for &x in v.iter() {
    new_v.push(x);
    new_v.push(x * x);
}
println!("{:?}", new_v);

но я хочу использовать итераторы. Я придумал этот код:

let v = vec![1, 2, 3];
let new_v: Vec<_> = v.iter()
    .flat_map(|&x| vec![x, x * x])
    .collect();
println!("{:?}", new_v);

но он выделяет промежуточный Vec в функции flat_map.

Как использовать flat_map без выделения?


iterator rust declarative
person Community    schedule 30.11.2019    source источник

Ответы (3)


arrow_upward
1
arrow_downward

Для этого можно использовать ArrayVec.

let v = vec![1, 2, 3];
let new_v: Vec<_> = v.iter()
    .flat_map(|&x| ArrayVec::from([x, x * x]))
    .collect();

Обсуждалось превращение массивов в итераторы по значению, чтобы вам не понадобилось ArrayVec, см. https://github.com/rust-lang/rust/issues/25725 и связанные PR.

person nnnmmm    schedule 30.11.2019

arrow_upward
2
arrow_downward

Начиная с версии 1.51.0 структура core::array::IntoIter была стабилизирована. Вы можете использовать это так:

use core::array;

let v = vec![1, 2, 3];
let new_v: Vec<_> = v.iter()
    .flat_map(|&x| array::IntoIter::new([x, x * x]))
    .collect();

Документация предупреждает, что это может быть устаревшим в будущем, когда IntoIterator будет реализовано для массивов, но в настоящее время это самый простой способ сделать это.

person Cheezey    schedule 02.04.2021
comment
Чтобы было ясно, this может быть устаревшим - это создание с помощью функции new, а не сам array::IntoIter. Если массивы когда-либо получат IntoIterator, вы сможете просто написать flat_map(|&x| [x, x * x]). - person trentcl; 03.04.2021

arrow_upward
1
arrow_downward

Если ваш итератор небольшой и вам не нужны внешние зависимости, короткий итератор можно создать из _ 1_ и _ 2_. Например,

use std::iter;

let v = vec![1, 2, 3];
let new_v: Vec<_> = v
    .iter()
    .flat_map(|&x| iter::once(x).chain(iter::once(x * x)))
    .collect();
println!("{:?}", new_v);

(игровая площадка)

Это можно было бы превратить в макрос, но имейте в виду, что использование этого для слишком большого количества элементов может привести к достижению предела рекурсии. Если вы делаете итератор для более чем нескольких десятков элементов, вероятно, не так уж плохо иметь выделение. Если вам действительно нужно небольшое увеличение производительности, решение nnnmmm, вероятно, будет лучше.

macro_rules! small_iter {
    () => { std::iter::empty() };
    ($x: expr) => {
        std::iter::once($x)
    };
    ($x: expr, $($y: tt)*) => {
        std::iter::once($x).chain(small_iter!($($y)*))
    };
}

fn main() {
    let v = vec![1, 2, 3];
    let new_v: Vec<_> = v
        .iter()
        .flat_map(|&x| small_iter!(x, x * x))
        .collect();
    println!("{:?}", new_v);
}

(игровая площадка)

person SCappella    schedule 30.11.2019