DecisionTree.cpp

// ============================================================
// Project 3: Decision Tree Classifier (ID3 Algorithm)
// OOP: Node, Tree, Dataset classes
// DSA: Recursion, BST-like tree splitting, entropy calculation
// Domain: AI/ML
// ============================================================

#include <iostream>
#include <vector>
#include <string>
#include <map>
#include <cmath>
#include <algorithm>
#include <iomanip>
#include <memory>
using namespace std;

// ---------- Dataset Row ----------
struct DataRow {
    vector<string> features;
    string label;
};

// ---------- Dataset ----------
class Dataset {
public:
    vector<string> featureNames;
    vector<string> labelName;
    vector<DataRow> rows;

    Dataset(const vector<string>& fNames) : featureNames(fNames) {}

    void addRow(const vector<string>& features, const string& label) {
        rows.push_back({features, label});
    }

    int numFeatures() const { return featureNames.size(); }

    // Count label occurrences
    map<string, int> labelCounts() const {
        map<string, int> counts;
        for (auto& row : rows) counts[row.label]++;
        return counts;
    }

    // Most common label
    string majorityLabel() const {
        auto counts = labelCounts();
        return max_element(counts.begin(), counts.end(),
            [](auto& a, auto& b){ return a.second < b.second; })->first;
    }

    // All same label?
    bool isPure() const {
        if (rows.empty()) return true;
        string first = rows[0].label;
        for (auto& r : rows)
            if (r.label != first) return false;
        return true;
    }

    // Get unique values for a feature
    vector<string> uniqueValues(int featureIdx) const {
        vector<string> vals;
        for (auto& r : rows) {
            if (find(vals.begin(), vals.end(), r.features[featureIdx]) == vals.end())
                vals.push_back(r.features[featureIdx]);
        }
        return vals;
    }

    // Split by feature value
    Dataset split(int featureIdx, const string& value) const {
        Dataset sub(featureNames);
        for (auto& r : rows)
            if (r.features[featureIdx] == value)
                sub.addRow(r.features, r.label);
        return sub;
    }
};

// ---------- Entropy & Information Gain ----------
class InfoGain {
public:
    static double entropy(const Dataset& ds) {
        if (ds.rows.empty()) return 0.0;
        auto counts = ds.labelCounts();
        double ent = 0.0;
        double total = ds.rows.size();
        for (auto& p : counts) {
            double prob = p.second / total;
            if (prob > 0) ent -= prob * log2(prob);
        }
        return ent;
    }

    static double informationGain(const Dataset& ds, int featureIdx) {
        double parentEntropy = entropy(ds);
        auto values = ds.uniqueValues(featureIdx);
        double weightedEntropy = 0.0;
        double total = ds.rows.size();

        for (auto& val : values) {
            Dataset sub = ds.split(featureIdx, val);
            weightedEntropy += (sub.rows.size() / total) * entropy(sub);
        }
        return parentEntropy - weightedEntropy;
    }

    static int bestFeature(const Dataset& ds, const vector<bool>& used) {
        int best = -1;
        double bestGain = -1.0;
        for (int i = 0; i < ds.numFeatures(); i++) {
            if (used[i]) continue;
            double gain = informationGain(ds, i);
            if (gain > bestGain) {
                bestGain = gain;
                best = i;
            }
        }
        return best;
    }
};

// ---------- Decision Tree Node ----------
struct TreeNode {
    bool isLeaf;
    string label;          // for leaf nodes
    int featureIdx;        // for internal nodes
    string featureName;
    map<string, shared_ptr<TreeNode>> children;

    TreeNode() : isLeaf(false), featureIdx(-1) {}
};

// ---------- Decision Tree ----------
class DecisionTree {
private:
    shared_ptr<TreeNode> root;

    shared_ptr<TreeNode> buildTree(const Dataset& ds, vector<bool> used, int depth) {
        auto node = make_shared<TreeNode>();

        // Base cases
        if (ds.rows.empty()) {
            node->isLeaf = true;
            node->label = "Unknown";
            return node;
        }
        if (ds.isPure()) {
            node->isLeaf = true;
            node->label = ds.rows[0].label;
            return node;
        }

        bool allUsed = true;
        for (bool u : used) if (!u) { allUsed = false; break; }
        if (allUsed || depth == 0) {
            node->isLeaf = true;
            node->label = ds.majorityLabel();
            return node;
        }

        // Find best split
        int bestFeat = InfoGain::bestFeature(ds, used);
        if (bestFeat == -1) {
            node->isLeaf = true;
            node->label = ds.majorityLabel();
            return node;
        }

        node->isLeaf = false;
        node->featureIdx = bestFeat;
        node->featureName = ds.featureNames[bestFeat];
        used[bestFeat] = true;

        for (auto& val : ds.uniqueValues(bestFeat)) {
            Dataset sub = ds.split(bestFeat, val);
            node->children[val] = buildTree(sub, used, depth - 1);
        }
        return node;
    }

    string predictNode(shared_ptr<TreeNode> node, const DataRow& row) const {
        if (node->isLeaf) return node->label;
        int fi = node->featureIdx;
        auto it = node->children.find(row.features[fi]);
        if (it == node->children.end()) return "Unknown";
        return predictNode(it->second, row);
    }

    void printTree(shared_ptr<TreeNode> node, const string& prefix, const string& edge) const {
        if (!edge.empty())
            cout << prefix.substr(0, prefix.size()-4) << "|-- [" << edge << "] ";
        if (node->isLeaf) {
            cout << "=> LABEL: " << node->label << "\n";
        } else {
            cout << "Feature: " << node->featureName << "\n";
            int i = 0;
            for (auto& child : node->children) {
                bool last = (++i == (int)node->children.size());
                printTree(child.second, prefix + (last ? "    " : "|   "), child.first);
            }
        }
    }

public:
    void train(const Dataset& ds, int maxDepth = 10) {
        vector<bool> used(ds.numFeatures(), false);
        root = buildTree(ds, used, maxDepth);
    }

    string predict(const DataRow& row) const {
        if (!root) return "Not trained";
        return predictNode(root, row);
    }

    void print() const {
        cout << "\n=== Decision Tree Structure ===\n";
        if (root) printTree(root, "", "");
    }

    double evaluate(const Dataset& ds) const {
        int correct = 0;
        for (auto& row : ds.rows)
            if (predict(row) == row.label) correct++;
        return (double)correct / ds.rows.size() * 100.0;
    }
};

// ---------- Main ----------
int main() {
    cout << "========================================\n";
    cout << "   Decision Tree Classifier (ID3)\n";
    cout << "========================================\n\n";

    // ---- Dataset 1: Weather -> Play Tennis ----
    cout << "--- Dataset: Play Tennis ---\n";
    Dataset tennis({"Outlook", "Temperature", "Humidity", "Wind"});

    tennis.addRow({"Sunny","Hot","High","Weak"},    "No");
    tennis.addRow({"Sunny","Hot","High","Strong"},  "No");
    tennis.addRow({"Overcast","Hot","High","Weak"}, "Yes");
    tennis.addRow({"Rain","Mild","High","Weak"},    "Yes");
    tennis.addRow({"Rain","Cool","Normal","Weak"},  "Yes");
    tennis.addRow({"Rain","Cool","Normal","Strong"},"No");
    tennis.addRow({"Overcast","Cool","Normal","Strong"},"Yes");
    tennis.addRow({"Sunny","Mild","High","Weak"},   "No");
    tennis.addRow({"Sunny","Cool","Normal","Weak"}, "Yes");
    tennis.addRow({"Rain","Mild","Normal","Weak"},  "Yes");
    tennis.addRow({"Sunny","Mild","Normal","Strong"},"Yes");
    tennis.addRow({"Overcast","Mild","High","Strong"},"Yes");
    tennis.addRow({"Overcast","Hot","Normal","Weak"},"Yes");
    tennis.addRow({"Rain","Mild","High","Strong"},  "No");

    DecisionTree dt;
    dt.train(tennis);
    dt.print();

    cout << "\n--- Predictions ---\n";
    vector<DataRow> testCases = {
        {{"Sunny","Cool","High","Weak"}, "?"},
        {{"Overcast","Hot","Normal","Strong"}, "?"},
        {{"Rain","Mild","High","Weak"}, "?"}
    };
    for (auto& tc : testCases) {
        cout << "  Outlook=" << tc.features[0]
             << " Temp=" << tc.features[1]
             << " Humidity=" << tc.features[2]
             << " Wind=" << tc.features[3]
             << " => Play: " << dt.predict(tc) << "\n";
    }

    double acc = dt.evaluate(tennis);
    cout << fixed << setprecision(1);
    cout << "\n  Training Accuracy: " << acc << "%\n";

    // ---- Dataset 2: Network Intrusion Detection (simplified) ----
    cout << "\n--- Dataset: Network Intrusion Detection ---\n";
    Dataset ids({"Protocol", "PacketSize", "Duration", "PortType"});

    ids.addRow({"TCP","Large","Long","Privileged"},   "Attack");
    ids.addRow({"UDP","Small","Short","Normal"},      "Normal");
    ids.addRow({"TCP","Large","Short","Privileged"},  "Attack");
    ids.addRow({"ICMP","Small","Long","Normal"},      "Attack");
    ids.addRow({"TCP","Small","Short","Normal"},      "Normal");
    ids.addRow({"UDP","Large","Long","Privileged"},   "Attack");
    ids.addRow({"TCP","Small","Long","Normal"},       "Normal");
    ids.addRow({"ICMP","Large","Short","Privileged"}, "Attack");
    ids.addRow({"UDP","Small","Short","Normal"},      "Normal");
    ids.addRow({"TCP","Large","Long","Normal"},       "Attack");

    DecisionTree ids_tree;
    ids_tree.train(ids);
    ids_tree.print();

    cout << "\n--- IDS Predictions ---\n";
    vector<DataRow> netTests = {
        {{"TCP","Large","Long","Privileged"}, "?"},
        {{"UDP","Small","Short","Normal"}, "?"},
        {{"ICMP","Small","Long","Normal"}, "?"}
    };
    for (auto& tc : netTests) {
        cout << "  Protocol=" << tc.features[0]
             << " Size=" << tc.features[1]
             << " Duration=" << tc.features[2]
             << " Port=" << tc.features[3]
             << " => " << ids_tree.predict(tc) << "\n";
    }

    cout << "\n  IDS Training Accuracy: " << ids_tree.evaluate(ids) << "%\n";
    return 0;
}