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.

Answer 2

I suppose, some kind of recursive DFS traversal algorithm should work

That's exactly right. Here is how you can construct this algorithm:

• Observe that the task has a recursive structure: applying it to any branch of a tree does to the branch the same thing that you want to do to the entire tree
• A branch could be either pruned, or removed entirely
• Your recursive implementation would return a pruned branch; it would signal branch removal by returning null
• The recursive function would check the incoming Node
• If the node represents a leaf, its content would be checked against the list of items that we wish to keep
• If a leaf is not on the "keep list", return null; otherwise return the leaf
• For non-leaf branches call the recursive method, and examine its result
• If the result is null, remove the corresponding branch from the map; otherwise, replace the branch with the pruned branch returned from the call
• If upon examining all branches the map of child nodes is empty, return null

Note that the algorithm could return null if none of the leaf nodes matches the "keep" list. If this is not desirable, add an extra level on top of the recursive implementation, and substitute null return at the top level with returning an empty tree.

Answer 3

With interface Tree:

public static interface Tree<T> {
    public T getValue();

    public List<Tree<T>> children();

    public default boolean isLeaf() {
        return children().isEmpty();
    }

    public default boolean removeDeadBranches(Predicate<T> testLiveLeaf) {
        if (isLeaf()) {
            return testLiveLeaf.test(getValue());
        }
        boolean remainLife = false;
        for (Iterator<Tree<T>> it = children().iterator(); it.hasNext();) {
            if (it.next().removeDeadBranches(testLiveLeaf)) {
                remainLife = true;
            } else {
                it.remove();
            }
        }
        return remainLife;
    }
}

Answer 4

import com.google.common.collect.Lists;
import org.junit.Before;
import org.junit.Test;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.util.List;

public class FilterTreeNode {
    private Logger logger = LoggerFactory.getLogger(FilterTreeNode.class);
    private TreeNode node0;
    private List<String> targetNode = Lists.newArrayList("B1","D1");

    @Before
    public void init(){
        node0 = TreeNode.builder().nodeCode("0").nodeName("A").build();
        TreeNode node1 = TreeNode.builder().nodeCode("1").nodeName("B").build();
        TreeNode node2 = TreeNode.builder().nodeCode("2").nodeName("C").build();
        TreeNode node3 = TreeNode.builder().nodeCode("3").nodeName("D").build();

        TreeNode node4 = TreeNode.builder().nodeCode("4").nodeName("B1").build();
        TreeNode node5 = TreeNode.builder().nodeCode("5").nodeName("B2").build();
        TreeNode node6 = TreeNode.builder().nodeCode("6").nodeName("C1").build();
        TreeNode node7 = TreeNode.builder().nodeCode("7").nodeName("D1").build();
        TreeNode node8 = TreeNode.builder().nodeCode("8").nodeName("D2").build();

        node1.setChildren(Lists.newArrayList(node4,node5));
        node2.setChildren(Lists.newArrayList(node6));
        node3.setChildren(Lists.newArrayList(node7,node8));

        node0.setChildren(Lists.newArrayList(node1,node2,node3));
    }

    @Test
    public void filterTest(){
        logger.info("before filter node0: {}",node0);
        List<TreeNode> retNodes = filterNode(node0);
        if (retNodes.size() >0){
            node0.setChildren(retNodes);
        }
        logger.info("after filter node0: {}",node0);
    }

    private List<TreeNode> filterNode(TreeNode treeNode){

        List<TreeNode> nodes = treeNode.getChildren();
        List<TreeNode> newNodes = Lists.newArrayList();
        List<TreeNode> tagNodes = Lists.newArrayList();

        for(TreeNode node : nodes){
            if (targetNode.contains(node.getNodeName())){
                newNodes.add(node);
            }
            if (node.getChildren() != null && node.getChildren().size() >0){
                List<TreeNode> retNodes = filterNode(node);
                if (retNodes.size()>0){
                    node.setChildren(retNodes);
                }else {
                    node.setChildren(null);
                    tagNodes.add(node);
                }
            }
        }
        nodes.removeAll(tagNodes);
        return newNodes;
    }
}

Answer 5

for any of you who look for a similar answer to @Kiril's for trees instead of map than:

the tree/node class:

@Getter
@Setter
@AllArgsConstructor
@NoArgsConstructor
@Builder
public class Tree {

    private TreeNode node;

    @Getter
    @Setter
    @AllArgsConstructor
    @NoArgsConstructor
    @Builder
    public static class TreeNode {
        private int id;
        private boolean isMatch;
        private List<TreeNode> nodes;
    }
}

and the filter method:

public Tree filter(Tree tree, List<Integer> ids) {
    if (tree.getNode() == null) {
        return tree;
    }
    filterNode(tree.getNode(), ids);

    return tree;
}

private boolean filterNode(Tree.TreeNode node, List<Integer> idsToFilter) {
    boolean isMatch = idsToFilter.contains(node.getId());
    node.setMatch(isMatch);
    if (CollectionUtils.isEmpty(node.getNodes()) && !isMatch) {
        return true;
    }

    node.getNodes().removeIf(treeNode -> filterNode(treeNode, idsToFilter));

    return node.getNodes().size() == 0 && !isMatch;
}