Saturday, May 14, 2016

Find all paths with specified sum in binary tree. (Java code)

package com.salpe;

import java.util.ArrayList;
import java.util.Collection;
import java.util.List;
import java.util.logging.Logger;

import org.junit.Test;

public class BSTTest2 {

 private static final Logger s_logger = Logger.getLogger(BSTTest2.class.getName());

 /**
  * Binary tree pre-order list as 0,1,2,3,4,5,6
  */
 @Test
 public void test() {

  // create test data
  BSTNode<Integer> node0 = new BSTNode<>(0, null, null);
  BSTNode<Integer> node2 = new BSTNode<>(2, null, null);
  
  BSTNode<Integer> node01 = new BSTNode<>(1, node0, node2);
  
  BSTNode<Integer> node06 = new BSTNode<>(6, null, null);
  BSTNode<Integer> node04 = new BSTNode<>(4, null, null);

  BSTNode<Integer> node05 = new BSTNode<>(5, node04, node06);
  
  BSTNode<Integer> node03 = new BSTNode<>(3, node01, node05);

  // create BST
  BST<Integer> bst = new BST<>(node03);

  Collection<List<Integer>> result = bst.findAllPaths(12);

  for (List<Integer> list : result) {

   s_logger.info(list.toString());

  }

 }

 /**
  * Binary Tree Node model
  *
  * @param <T>
  */
 static class BSTNode<T extends Comparable<T>> {

  T data;
  BSTNode<T> left;
  BSTNode<T> right;

  BSTNode(T data, BSTNode<T> left, BSTNode<T> right) {

   this.data = data;
   this.left = left;
   this.right = right;

  }

  @Override
  public String toString() {
   return "BSTNode [data=" + data + "]";
  }

 }

 /**
  * Binary Tree Representation
  *
  * @param <T>
  */
 static class BST<T extends Comparable<T>> {

  BSTNode<T> root;

  public BST(BSTNode<T> root) {
   this.root = root;
  }

  public Collection<List<T>> findAllPaths(int sum) {

   List<T> paths = new ArrayList<>();
   Collection<List<T>> result = new ArrayList<>();
   depthFirstSearch0(root, 0, sum, paths, result);
   return result;

  }

  public void depthFirstSearch0(BSTNode<T> node, int currentSum, int expectedSum, List<T> paths,
    Collection<List<T>> result) {

   if (node == null) {

    return;
   }

   paths.add(node.data);

   // check sum
   currentSum = currentSum + Integer.parseInt(node.data.toString());

   if (currentSum == expectedSum) {
    // add cloned list
    result.add(new ArrayList<>(paths));
   }
            // go left 
   depthFirstSearch0(node.left, currentSum, expectedSum, paths, result);
    // then go right
   depthFirstSearch0(node.right, currentSum, expectedSum, paths, result);

   // remove when going up in level
   paths.remove(paths.size() - 1);

  }

 }

}

No comments: