CODEWITHSUNDEEP
nestedrust
Geodesic
Geodesic

Reputation: 893

Transforming a list of structs with parent IDs into a list of trees

I'm pulling a dataset from a database which gives me a vector of structs of the form:

struct Foo {
    id: i32,
    parent: Option<i32>,
    data: String,
}

I would like to serialize and output to JSON the nested version of this data as a vector of:

struct Bar {
    id: i32,
    data: String,
    children: Option<Vec<Bar>>,
}

I'm having some issues wrapping my head around the implementation of this due to the recursive nature. I can solve the problem one level down using iterators, but then hit a wall when I want to iterate again over the same vector.

For example, a method on Vec<Foo> which attempts to just nest children ids into a hashmap:

fn build_tree(&self) -> HashMap<i32, Vec<i32>> {
    let mut tree = HashMap::new();
    for node in self.iter() {
        if let Some(parent) = node.parent {
            let leaf = tree.entry(parent).or_insert(Vec::new());
            leaf.push(node.id);
        }
    }
    tree
}

yields

{14: [15], 3: [14], 1: [2, 17], 2: [16, 18], 18: [19], 19: [20]}

But what I require would be something deeper:

{3: [14: [15]], 1: [2: [16, 18: [19: [20]]], 17]}

Reading through this post about turning a recursive idea into iterative code suggests that such an implementation is possible, but I've had difficulty taking the ideas from that problem and applying them here.

Can someone describe a method for transforming this Vec<Foo> to a Vec<Bar>? I'd be happy with an iterative or recursive suggestion; I just had a lot of issues with borrows and referencing when I tried recursion myself.

Upvotes: 2

Views: 226

Answers (1)

Shepmaster
Shepmaster

Reputation: 430673

The straight-line solution involves building up a graph of all your data and traversing it recursively, returning the Bar from each level and collecting them.

We start by creating a petgraph::DiGraphMap — a directed graph that allows us to control the node IDs (since we just have numeric identifiers). If a node has a parent, we ensure that exists in the graph and add an edge from the parent to the child. If it doesn't have a parent, we know that it will be one of our top-level ids, so we stash it aside for later:

let mut graph = DiGraphMap::new();
let mut top_level_ids = vec![];

for i in &input {
    graph.add_node(i.id);

    match i.parent {
        Some(parent_id) => {
            graph.add_node(parent_id);
            graph.add_edge(parent_id, i.id, ());
        }
        None => {
            top_level_ids.push(i.id);
        }
    }
}

Next, we iterate over all of the top-level IDs and convert them into a Bar:

let result: Vec<_> = top_level_ids
    .into_iter()
    .map(|id| build_tree(&graph, id))
    .collect();

Building a Bar is the recursive core of the problem. We construct another Bar for each child, stuff them all into a Vec, then return the current Bar:

fn build_tree(graph: &DiGraphMap<i32, ()>, id: i32) -> Bar {
    let children = graph
        .neighbors(id)
        .map(|child_id| build_tree(graph, child_id))
        .collect();

    Bar { id, children }
}

At this point, you have a Vec<Bar>. It's an exercise for the reader how to properly encode this for the desired JSON format :-).

The complete example.

Upvotes: 3

Related Questions