Find Sum of all Levels in a Binary Search Tree

package org.wesome.dsalgo;//package org.wesome.dsalgo;
import lombok.Data;

import java.util.Objects;

@Data
class Node {
    int data;
    Node left, right;

    public Node(int data) {
        this.data = data;
        left = null;
        right = null;
    }
}

public class BinaryTree {
    Node root;

    void insert(int data) {
        System.out.println("BinarySearchTree.insert data = " + data);
        root = insert(root, data);
    }

    Node insert(Node root, int data) {
        if (Objects.isNull(root)) {
            Node tempNode = new Node(data);
            return tempNode;
        }
        if (data > root.data) {
            root.right = insert(root.right, data);
        } else {
            root.left = insert(root.left, data);
        }
        return root;
    }

    public static int[] getSumOfLevel(Node root, int level, int[] arr) {
        if (Objects.isNull(root)) {
            return arr;
        }
        arr[level] = arr[level] + root.data;
        getSumOfLevel(root.left, level + 1, arr);
        getSumOfLevel(root.right, level + 1, arr);
        return arr;
    }

    /*  In order to calculate the sum of each level, we need the height of binary tree  */
    static int calculateHeight(Node root) {
        if (Objects.isNull(root.left) && Objects.isNull(root.right)) {
            System.out.println("left and right of " + root.data + " is null, hence returning 0");
            return 0;
        }
        int leftHeight = 0;
        if (Objects.nonNull(root.left)) {
            leftHeight = calculateHeight(root.left);
        }
        int rightHeight = 0;
        if (Objects.nonNull(root.right)) {
            rightHeight = calculateHeight(root.right);
        }
        return (Math.max(leftHeight, rightHeight) + 1);
    }

    private static void print(int[] sumOfLevel) {
        for (int i = 0; i < sumOfLevel.length; i++) {
            System.out.println(i + " level sum is  = " + sumOfLevel[i]);
        }
    }

    /*  height of binary tree will be the size of array, and sum will be stored in array with level as index */
    public int[] findSumOfEachLevel(Node root) {
        int height = calculateHeight(root) + 1;
        int[] arr = new int[height];
        int[] sumOfLevel = getSumOfLevel(root, 0, arr);
        print(sumOfLevel);
        return sumOfLevel;
    }
}
package org.wesome.dsalgo;

import org.junit.jupiter.api.Assertions;
import org.junit.jupiter.api.Test;

public class BinaryTreeTest {
    @Test
    void findSumOfEachLevelTest() {
        BinaryTree binaryTree = new BinaryTree();
        binaryTree.insert(4);
        binaryTree.insert(2);
        binaryTree.insert(1);
        binaryTree.insert(3);
        binaryTree.insert(6);
        binaryTree.insert(5);
        binaryTree.insert(7);
        binaryTree.insert(7);
        Assertions.assertArrayEquals(new int[]{4, 8, 16, 7}, binaryTree.findSumOfEachLevel(binaryTree.root));
    }
}
plugins {
    id 'java'
    id "io.freefair.lombok" version "6.2.0"
}

group = 'org.wesome'
version = '0.0.1-SNAPSHOT'
sourceCompatibility = JavaVersion.VERSION_1_8

configurations {
    compileOnly {
        extendsFrom annotationProcessor
    }
}

repositories {
    mavenCentral()
}

dependencies {
    testImplementation 'org.junit.jupiter:junit-jupiter:5.6.2'
}

test {
    useJUnitPlatform()
}

 

follow us on