如何用rust写一个把对象数组转成树结构的通用方法

alexnest 发表于 2023-03-08 15:21

最近用 rust 写 web 框架，无论菜单还是组织架构都需要把对象数组转成树结构，希望抽取出一个通用方法，以下代码可以跑起来，但效率极低，因为无脑 clone，也尝试用智能指针去优化本段代码，但鉴于所有权和生命周期的报错和本人能力有限，未能编译通过。希望群友可以给出一个思路去优化本段代码或者提供一个执行效率比较高的通用方法。

serde = { version = "1.0", features = ["derive"] }

use std::collections::HashMap;
use std::fmt::Debug;
use std::hash::Hash;

#[derive(Debug, Clone, Default)]
struct TreeNode<T>
    where T: Clone + Default + Debug
{
    item: T,
    children: Option<Vec<TreeNode<T>>>,
}

impl<T> TreeNode<T>
    where T: Clone + Default + Debug
{
    fn init_tree(items: Vec<T>) -> Vec<TreeNode<T>> {
        let mut tree_items: Vec<TreeNode<T>> = Vec::new();
        for mut item in items {
            let tree = TreeNode { item, ..Default::default() };
            tree_items.push(tree);
        }
        tree_items
    }

    fn set_tree<F1, F2>(tree_items: Vec<TreeNode<T>>, parent_key: String, get_key: F1, get_parent_key: F2) -> Vec<TreeNode<T>>
        where F1: Fn(&T) -> String + Clone,
              F2: Fn(&T) -> String + Clone
    {
        let mut trees: Vec<TreeNode<T>> = Vec::new();
        for mut tree in tree_items.clone() {
            if get_parent_key(&tree.item) == parent_key {
                tree.children = Some(Self::set_tree(tree_items.clone(), get_key(&tree.item), get_key.clone(), get_parent_key.clone()));
                trees.push(tree.clone())
            }
        }
        trees
    }

    fn build_tree<F1,F2>(items: Vec<T>, parent_key: String, get_key: F1, get_parent_key: F2) -> Vec<TreeNode<T>>
        where F1: Fn(&T) -> String + Clone,
              F2: Fn(&T) -> String + Clone
    {
        let items = Self::init_tree(items);
        Self::set_tree(items,parent_key,get_key,get_parent_key)
    }
}


fn main() {
    #[derive(Debug, Clone, Default)]
    struct Item {
        id: String,
        parent_id: Option<String>,
        name: String,
    }

    let items = vec![
        Item {
            id: String::from("1"),
            parent_id: Some(String::from("0")),
            name: String::from("Root"),
        },
        Item {
            id: String::from("2"),
            parent_id: Some(String::from("1")),
            name: String::from("Child 1"),
        },
        Item {
            id: String::from("3"),
            parent_id: Some(String::from("1")),
            name: String::from("Child 2"),
        },
        Item {
            id: String::from("4"),
            parent_id: Some(String::from("2")),
            name: String::from("Grandchild 1.1"),
        },
        Item {
            id: String::from("5"),
            parent_id: Some(String::from("2")),
            name: String::from("Grandchild 1.2"),
        },
    ];

    let tree = TreeNode::build_tree(items,String::from("0"),|i|i.id.clone(),|i|i.parent_id.clone().unwrap().clone());
    println!("tree is {:?}",tree);
}

评论区

写评论

asuper 2023-05-17 11:31

如果树不需要修改或改动少，还有一种避免clone的办法，另外开一个索引数组，在数组上实现链表，但只存下标不存实际元素

作者 alexnest 2023-03-09 16:21

用以下测试用例比较了这种写法和另外几种写法，这种写法的运行效率是最高的。我的那种写法时间复杂度和空间复杂度都是最差的.thx a lot!

println!("init data start!");
let mut items = Vec::new();
for i in 1..=100000 {
    let key = i.to_string();
    let mut rng = rand::thread_rng();
    let rand_num = rng.gen_range(0..100);
    let parent_key = rand_num.to_string();

    let item = Item {
        id: key.clone(),
        parent_id: Some(parent_key.clone()),
        name: format!("{}.{}",&key,&parent_key)
    };
    items.push(item);
}
println!("init data complete!");
let items1 = items.clone();
println!("clone data complete!");

let start_time = SystemTime::now();
let t1 = tree_node_2::TreeNode::build_tree(items, String::from("0"), |i| i.id.clone(), |i| i.parent_id.clone().unwrap().clone());
let end_time = SystemTime::now();
let duration = end_time.duration_since(start_time).unwrap();
let duration_seconds = duration.as_secs();
println!("Time elapsed: {} seconds", duration_seconds);

let start_time = SystemTime::now();
items1.to_tree(Some(0.to_string()));
let end_time = SystemTime::now();
let duration = end_time.duration_since(start_time).unwrap();
let duration_seconds = duration.as_secs();
println!("Time elapsed: {} seconds", duration_seconds);

--
👇
4j3n: ```rust use std::fmt::Debug; #[derive(Debug, Clone, Default)] struct TreeNode where T: Clone + Default + Debug, { item: T, children: Option<Vec<TreeNode>>, //XXX: 为啥不直接Vec? }

impl TreeNode where T: Clone + Default + Debug, { fn build_tree<F1, F2>( items: Vec, root: String, get_key: F1, get_parent_key: F2, ) -> Vec<TreeNode> where F1: Fn(&T) -> String + Clone, F2: Fn(&T) -> String + Clone, { let mut map = std::collections::HashMap::new(); for item in items { map.entry(get_parent_key(&item)) .or_insert_with(Vec::new) .push(TreeNode { item, children: None, }); } Self::build(&mut map, root, &get_key).unwrap_or(Vec::new()) }

fn build(
    map: &mut std::collections::HashMap<String, Vec<TreeNode<T>>>,
    root: String,
    key: &impl Fn(&T) -> String,
) -> Option<Vec<TreeNode<T>>> {
	//map.get(&root).clone() if duplicate key
    map.remove(&root).map(|mut children| {
        for child in &mut children {
            child.children = Self::build(map, key(&child.item), key);
        }
        children
    })
}

}

原来写法也忒离谱了，dfs代码量也没有更多；就算不用rust，其他语言也不能这么写啊。。。。。。

‘static 2023-03-09 10:42

需要复制指针吗？

#[derive(Debug)]
struct TreeNodePtr<'a, T> {
    item: &'a T,
    children: Option<Vec<TreeNodePtr<'a, T>>>,
}

trait Tree {
    type Key: std::cmp::PartialEq;
    fn get_key(&self) -> Self::Key;
    fn get_parent_key(&self) -> Self::Key;
}

trait BuildTree {
    type Item: Tree;
    fn to_tree(
        &self,
        parent_key: <Self::Item as Tree>::Key,
    ) -> Option<Vec<TreeNodePtr<'_, Self::Item>>>;
}

impl<T, K> BuildTree for Vec<T>
where
    K: std::cmp::PartialEq,
    T: Tree<Key = K>,
{
    type Item = T;
    fn to_tree(
        &self,
        parent_key: <Self::Item as Tree>::Key,
    ) -> Option<Vec<TreeNodePtr<'_, Self::Item>>> {
        let mut tree = Vec::new();
        for item in self {
            if item.get_parent_key() == parent_key {
                let new_parent_key = item.get_key();
                let children = self.to_tree(new_parent_key);
                let node = TreeNodePtr { item, children };
                tree.push(node);
            }
        }
        if tree.is_empty() {
            None
        } else {
            Some(tree)
        }
    }
}
fn main(){
    #[derive(Debug, Clone, Default)]
    struct Item {
        id: String,
        parent_id: Option<String>,
        name: String,
    }
    
    impl Tree for Item {
        type Key = Option<String>;
        fn get_key(&self) -> Self::Key {
            Some(self.id.clone())
        }
        fn get_parent_key(&self) -> Self::Key {
            self.parent_id.clone()
        }
    }
    
    let items = vec![
        Item {
            id: String::from("1"),
            parent_id: Some(String::from("0")),
            name: String::from("Root"),
        },
        Item {
            id: String::from("2"),
            parent_id: Some(String::from("1")),
            name: String::from("Child 1"),
        },
        Item {
            id: String::from("3"),
            parent_id: Some(String::from("1")),
            name: String::from("Child 2"),
        },
        Item {
            id: String::from("4"),
            parent_id: Some(String::from("2")),
            name: String::from("Grandchild 1.1"),
        },
        Item {
            id: String::from("5"),
            parent_id: Some(String::from("2")),
            name: String::from("Grandchild 1.2"),
        },
    ];
let tree = items.to_tree(Some(0.to_string()));

    println!("tree is {:?}", tree);
}

4j3n 2023-03-09 00:51

use std::fmt::Debug;
#[derive(Debug, Clone, Default)]
struct TreeNode<T>
where
    T: Clone + Default + Debug,
{
    item: T,
    children: Option<Vec<TreeNode<T>>>, //XXX: 为啥不直接Vec?
}

impl<T> TreeNode<T>
where
    T: Clone + Default + Debug,
{
    fn build_tree<F1, F2>(
        items: Vec<T>,
        root: String,
        get_key: F1,
        get_parent_key: F2,
    ) -> Vec<TreeNode<T>>
    where
        F1: Fn(&T) -> String + Clone,
        F2: Fn(&T) -> String + Clone,
    {
        let mut map = std::collections::HashMap::new();
        for item in items {
            map.entry(get_parent_key(&item))
                .or_insert_with(Vec::new)
                .push(TreeNode {
                    item,
                    children: None,
                });
        }
        Self::build(&mut map, root, &get_key).unwrap_or(Vec::new())
    }

    fn build(
        map: &mut std::collections::HashMap<String, Vec<TreeNode<T>>>,
        root: String,
        key: &impl Fn(&T) -> String,
    ) -> Option<Vec<TreeNode<T>>> {
		//map.get(&root).clone() if duplicate key
        map.remove(&root).map(|mut children| {
            for child in &mut children {
                child.children = Self::build(map, key(&child.item), key);
            }
            children
        })
    }
}

原来写法也忒离谱了，dfs代码量也没有更多；就算不用rust，其他语言也不能这么写啊。。。。。。

ifyuhj 2023-03-08 20:44

按你的思路改了一下，避免了不必要的复制，但是每个Item肯定得复制一次的。

fn build_tree2<F1, F2>(
    items: &[T],  
    parent_key: &str,
    get_key: &F1,
    get_parent_key: &F2,
) -> Vec<TreeNode<T>>
where
    F1: Fn(&T) -> String + Clone,
    F2: Fn(&T) -> String + Clone,
{
    let mut trees: Vec<TreeNode<T>> = Vec::new();
    for item in items.iter() {
        if get_parent_key(item).eq(parent_key) {
            let mut node = TreeNode {
                item: item.clone(),
                children: None,
            };
            let parent_key = get_key(item);
            node.children = Some(Self::build_tree2(
                items,
                &parent_key,
                get_key,
                get_parent_key,
            ));
            trees.push(node);
        }
    }
    trees
}

songzhi 2023-03-08 17:40

像一楼说的，ID用满足Copy的类型最好。如果不想Clone Item，可以先用key构建出一棵树出来，再从root_key深度优先遍历一次，就能不Clone只move构建出整颗树。如果不想这么麻烦，那就用智能指针, 但是用了智能指针之后TreeNode的类型定义会很丑。

wangbyby 2023-03-08 16:11

String作为ID不是很好，ID是usize或许更好一点
set_tree的逻辑是O(n^2)的，当然会慢
如果可能的话，对数组按照parent_id排序然后建立树会更快一点

1 共 7 条评论, 1 页