general_tree.py

# -*- coding: utf-8 -*-
"""datastructure_tree.ipynb

Automatically generated by Colaboratory.

Original file is located at
    https://colab.research.google.com/drive/1Of-vTd0-JTwoT3UAazDxZgu35rZrVEYJ
"""

class TreeNode:
  def __init__(self, data):
    self.data = data
    self.children = []
    self.parent = None

  def add_child(self, child):
    child.parent = self
    self.children.append(child)

  def get_parent(self):
    level = 0
    p = self.parent
    while p:
      level += 1
      p = p.parent
    
    return level

  def display(self):
    spaces = " " * self.get_parent() * 3
    prefix = spaces + "|__" if self.parent else ""
    print(prefix + self.data)
    if self.children:
      for child in self.children:
        child.display()


def main_tree():
  # define root
  root = TreeNode("Machine Learning Algorithms")

  # define child level 1
  sup = TreeNode("Supervised Learning")
  sup.add_child(TreeNode("Linear Regression"))
  sup.add_child(TreeNode("Logistic Regression"))
  sup.add_child(TreeNode("Decision Tree"))
  sup.add_child(TreeNode("Support Vector Machine"))

  unsup = TreeNode("Unsupervised Learning")
  unsup.add_child(TreeNode("Hierarchical Clustering"))
  unsup.add_child(TreeNode("K-Means Clustering"))
  unsup.add_child(TreeNode("K-Nearest Neighbors"))

  # adding a root child
  root.add_child(sup)
  root.add_child(unsup)

  # print the entire root
  root.display()

if __name__ == '__main__':
  main_tree()