DecisionTree.py

from collections import defaultdict
from Node import Node

class DecisionTree:
    # initializes decision tree
    def __init__(self, data: list, min_sample_split: int, max_depth = 10):
        self.max_depth = max_depth
        self.min_sample_split = min_sample_split
        self.root = Node(data)

    # caller function to recursively build tree
    def build(self):
        self.root = self.build_helper(self.root, 0)      
    
    # recursive function to construct tree
    def build_helper(self, n: Node, cur_depth: int):
        # base case
        if (cur_depth == self.max_depth or len(n.data) < self.min_sample_split or self.is_pure(n.data)):
            n.pred = self.calc_majority(n.data)
            return n

        # recursive case
        # select the best split at this node
        partitions = self.select_best_split(n)

        # check valid partitions
        if (partitions is None or len(partitions[0]) == 0 or len(partitions[1]) == 0):
            # there is no valid partition which means that this is a leaf node
            n.pred = self.calc_majority(n.data)
            return n
    
        n.left = self.build_helper(Node(partitions[0]), cur_depth + 1)
        n.right = self.build_helper(Node(partitions[1]), cur_depth + 1)
        
        return n

    # caller function to recursively classify guess
    def predict(self, guess: list):
        return self.predict_help(guess, self.root)
    
    # recursive function to navigate tree and classify guess
    def predict_helper(self, guess: list, n: Node):
        # base case: leaf node
        if (n.is_leaf()):
            return n.pred
        
        # recursive case: internal node
        if (guess[n.feature_index] <= n.split_value):
            result = self.predict_helper(guess, n.left)
        else:
            result = self.predict_helper(guess, n.right)
        
        return result

    # returns true if all samples in data belong to only one class, false otherwise
    def is_pure(self, data: list):
        chosen = data[0][-1]

        # check the class each sample against chosen
        for i in range(1, len(data)):
            if (data[i][-1] != chosen):
                return False
            
        return True

    # calculates the majority label in data
    def calc_majority(self, data: list):
        occ_dict = defaultdict(int)

        maj_label = None
        maj_occurences = 0

        # tracks occurences of each sample in data and majority label
        for sample in data:
            label = sample[-1]
            occ_dict[label] += 1
            if (occ_dict[label] > maj_occurences):
                maj_occurences = occ_dict[label]
                maj_label = label
        
        return maj_label
    
    # select the best split of data for the given node
    def select_best_split(self, n: Node):
        labels = self.extract_labels(n.data)

        n_features = len(n.data[0]) - 1

        best_groups = None

        # check every possible feature
        for feature_index in range(n_features):
            unique_values = set(row[feature_index] for row in n.data)

            # check every unique splitting value to test all possible splits
            for splitting_value in unique_values:
                groups = self.get_split(feature_index, splitting_value, n.data)
                cur_gini = self.gini(groups, labels)
                if (cur_gini < n.gini):
                    best_groups = groups
                    n.set_split(feature_index, splitting_value, cur_gini)
        return best_groups

    # extract labels from data
    def extract_labels(self, data: list) -> list:
        labels = set()

        # add all label possibilities to 'labels' set
        for sample in data:
            labels.add(sample[-1])
        
        return list(labels)

    # splits the data into left and right child nodes based on a 
    # input feature index and threshold value
    def get_split(self, index: int, value: int, data: list) -> tuple[list, list]:
        # index: index of the feature to use to create the split
        # value: the threshold value that decides left/right child
        # data: the dataset

        left, right = [], []

        # checks each sample and partitions accordingly
        for sample in data:
            if (sample[index] <= value):
                left.append(sample)
            else:
                right.append(sample)
        return [left, right]

    # calculates the gini score for a particular split
    def gini(self, groups: list, y: list) -> float:
        # y is the class labels
        # groups represents the 2 child nodes for this particular split
        
        total_samples = 0.0

        # count total samples
        for child in groups:
            total_samples += len(child)

        gini = 0.0 # total gini

        # groups is already divided by the split. now we need to go in each group
        for child in groups:
            gini_group = 0.0 # this is the group specific gini (1 - sum(prop_i^2))
            
            if (len(child) == 0):
                # this post-split group has no samples
                continue

            # calculate proportion of samples belonging to a each class
            for class_label in y:
                # now we need to see how many samples in child belong to class_label
                gini_group += self.get_proportion(child, class_label) ** 2

            gini += (1 - gini_group) * (len(child) / total_samples)
        
        return gini

    # calculates the proportion of samples within child that are 
    # classified under cur_label
    def get_proportion(self, child: list, cur_label: str) -> float:
        # child is the current child node data that we are dealing with
        # cur_label is the current label we are looking at
        label_count = 0.0

        # count number of samples in this child node that are classified as cur_label
        for sample in child:
            if (sample[-1] == cur_label):
                label_count += 1
        
        return label_count / len(child)
    
    # caller function to recursively print this tree
    def print_tree(self):
        self.print_tree_helper(self.root, "")

    # recursive function to print this tree
    def print_tree_helper(self, n: Node, spaces: str):
        if (n is not None):
            self.print_tree_helper(n.right, spaces + "    ")
            print("\n" + spaces + n.get_str_data())
            self.print_tree_helper(n.left, spaces + "    ")