Javascript: How to preserve tree structure and grandparent-child relationship when user selects/deselects nodes from the tree

Question

I have a tree structure {name, [children]} (JS object). The user can arbitrarily select any nodes in order to partially duplicate the structure (into an array of JS objects, since loose leaf nodes could be selected). eg the user could select a leaf node but not its parent/grandparent, then the leaf node will just be in the flat array in the children of the highest parent.

Example:

original (JS Object)      selected
   A✓                       A
   |  \ \                  / \
   B✓ C D✓                B  D
  / \          =>         / \
 E✓  F                   E  G
     |
     G✓ 

original                  selected(Array of JS objects)
   A                       B, D      
   |  \ \                 / \
   B✓ C D✓               E  G 
  / \          =>         
 E✓  F                   
     |
     G✓ 

My first thoughts were: for each node on click, assume we maintain a selected array:

1. if this node is in selected => deselect (fairly straight forward):

   a1. if the node is a "top-level" element => remove it from the array, and push its children to the array.

   a2. if the node is a (grand)child of a top-level element => remove the node from its parent's children, and reattach the node's children to its parent.

2. if this node is not in selected => select (weird things happen):

   a. Initialize a tempNode with an empty children , for we don't know its children yet

   b. if this node has (grand)children => for each (grand)child, if the (grand)child is in selected => remove the (grand)child in selected => add this (grand)child to the tempNode 's children (should be recursive for multiple levels)

   c1. if this node's (grand)parent is in selected already => attach the tempNode to the (grand)parent (should be recursive for multiple levels)

   c2. if this node's (grand)parent is not in selected , push it to the array.

I'm mainly stuck on step 2b and 2c.

//for adding node into selected
const addChildInSelected = (child) => {
    var tempNode = {
        name: child.name,
        children: []
    }
    var clonedState = clone(selected)
    var childInTree = findChildInTree(child.id, props.data)
    //if the child has children, check if they are in selected yet.
    if (childInTree.children.length > 0) {
        //Problem 1: I should run this recursively for all the grandchildren, 
        //but I struggle to find the base case. But recursion might not be the best solution
        for (var i = 0; i < childInTree.children.length; i++) {
            //Problem 2: findChildInSelected/removeChildInSelect are all recursive calls as well,
            //due to the tree structure. very worried about stack overflow during this step...
            var grandChildInSelected = findChildInSelected(childInTree.children[i].id)
            if (grandChildInSelected !== null) {
                clonedState = removeChildInSelected(clonedState)
                tempNode.children.push(findChildInTree(childInTree.children[i]))
            }
        }
    }
    //Problem 3. I realized I needed to check for each node in the `selected` again. 
    //another potential performance hurdle
}

Now that I rethink about this issue, each node only cares about itself, its immediate parent/grandparent, and all of its children/grandchildren as a flat array. That might be a good perspective to look into this problem. Maybe a helper function that finds the parent/grandparent or all children would be a good start. (Plus I realized that in several functions I forgot to check grandparent-child...)

But since I'm operating with multiple tree structures, the function below just doesn't seem ideal...

//just a quick example
const findDescendantsInSelected = (node, descendants=[]) => {
    //recursive call within recursive function which could overflow if the tree is big
    if (findNodeInSelected(node) !== null) {
        descendents.push(node)
    }
    if (node.children.length > 0) {
        for (var i = 0; i < entry.children.length; i++) {
            result = findDescendantsInSelected(node.children[i], descendants);
        }
        return result
    }
    return descendents
}

All in all, maintaining almost 3 separate trees (selected state, original data, and clone on each step since I'm using react) and tracking the parents/children in each of them, across different functions made my head hurt. Any insight on how to simplify the problem would be greatly appreciated!

Answer 1

A simple and pure recursive function that uses a depth-first traversal should suffice. It returns the list of selected trees you want (known as a forest ). If a node is selected, it will collect all the selected descendant trees and return a new node that has these as children. If the node is not selected, it just joins the lists from its children and returns that list.

function selectedForest(node) {
    const trees = node.children.flatMap(selectedForest);
    return isSelected(node) ? [ {...node, children: trees} ] : trees;
}

Demo:

 const demo = { name: "A", children: [ { name: "B", children: [ { name: "E", children: [] }, { name: "F", children: [ { name: "G", children: [] }, ] }, ] }, { name: "C", children: [] }, { name: "D", children: [] }, ] }; let selection; function isSelected({name}) { return selection.has(name); } function selectedForest(node) { const trees = node.children.flatMap(selectedForest); return isSelected(node)? [ {...node, children: trees} ]: trees; } selection = new Set("ABDEG"); console.log(JSON.stringify(selectedForest(demo), null, 2)); selection = new Set("BDEG"); console.log(JSON.stringify(selectedForest(demo), null, 2));