Java generic tree traversal with node filtering

Question

I have a generic tree structure. I need an algorithm to traverse it, and remove some leafs if they are not contained in a given list. If all the leafs are removed from a subtree, then remove the whole subtree too.

Example tree:

         0
       / |  \
      1  2   3
    /  \ |  / \
   4   5 6  7  8

Leafs to keep: {4, 6}

Result tree:

         0
       / | 
      1  2 
    /    | 
   4     6  

The input data structure is contained in a HashMap where the key is the parent id of the nodes, and the value is a List of the Nodes directly under the parent (but not recursively all children). The parent id of the root node is empty string.

Map<String, List<Node>> nodes;

class Node {
    private String id;
    //other members
    //getters, setters
}

I suppose, some kind of recursive DFS traversal algorithm should work, but I couldn't find it out, how.

Answer 1

I suggest you to try the following approach:

The method boolean removeRecursively(String id, Set<String> leavesToKeep) will traverse from the node with the given id down to this branch leaves.

First we check whether the curent node is a leaf or not. If the leaf is not in the leavesToKeep set, we remove it and return true, otherwise return false. This is a base case of our recursion.

If the node is not a leaf, then we do something like this:

children.removeIf(n -> removeRecursively(n.id, leavesToKeep));

removeIf is a convenient Java 8 method to remove all of the elements that satisfy the given predicate. This means that the child will be removed from the list only if all of its children are removed too. Therefore, we should make removeRecursively return true if after the children.removeIf call the children list is empty:

if (children.isEmpty()) {
    tree.remove(id);
    return true;
} else return false;

Full method may look like this:

public static boolean removeRecursively(Map<String, List<Node>> tree, String id, Set<String> leavesToKeep) {
    List<Node> children = tree.get(id);
    if (children == null || children.isEmpty()) {
        if (!leavesToKeep.contains(id)) {
            tree.remove(id);
            return true;
        } else return false;
    }
    children.removeIf(n -> removeRecursively(tree, n.id, leavesToKeep));
    if (children.isEmpty()) {
        tree.remove(id);
        return true;
    } else return false;
}

where tree is the map you described, id is the start node id, leavesToKeep is a set of ids to keep.