Add cycle detect adding edges

Network.cc

@@ -6,30 +6,75 @@ namespace bayesnet {
             delete pair.second;
         }
     }
-    void Network::addNode(std::string name, int numStates)
+    void Network::addNode(string name, int numStates)
     {
         nodes[name] = new Node(name, numStates);
+        if (root == nullptr) {
+            root = nodes[name];
+        }
     }
-    void Network::addEdge(const std::string parent, const std::string child)
+    void Network::setRoot(string name)
+    {
+        if (nodes.find(name) == nodes.end()) {
+            throw invalid_argument("Node " + name + " does not exist");
+        }
+        root = nodes[name];
+    }
+    Node* Network::getRoot()
+    {
+        return root;
+    }
+    bool Network::isCyclic(const string& nodeId, unordered_set<string>& visited, unordered_set<string>& recStack)
+    {
+        if (visited.find(nodeId) == visited.end()) // if node hasn't been visited yet
+        {
+            visited.insert(nodeId);
+            recStack.insert(nodeId);
+
+            for (Node* child : nodes[nodeId]->getChildren()) {
+                if (visited.find(child->getName()) == visited.end() && isCyclic(child->getName(), visited, recStack))
+                    return true;
+                else if (recStack.find(child->getName()) != recStack.end())
+                    return true;
+            }
+        }
+        recStack.erase(nodeId); // remove node from recursion stack before function ends
+        return false;
+    }
+    void Network::addEdge(const string parent, const string child)
     {
         if (nodes.find(parent) == nodes.end()) {
-            throw std::invalid_argument("Parent node " + parent + " does not exist");
+            throw invalid_argument("Parent node " + parent + " does not exist");
         }
         if (nodes.find(child) == nodes.end()) {
-            throw std::invalid_argument("Child node " + child + " does not exist");
+            throw invalid_argument("Child node " + child + " does not exist");
         }
+        // Temporarily add edge to check for cycles
         nodes[parent]->addChild(nodes[child]);
         nodes[child]->addParent(nodes[parent]);
+        // temporarily add edge
+
+        unordered_set<string> visited;
+        unordered_set<string> recStack;
+
+        if (isCyclic(nodes[child]->getName(), visited, recStack)) // if adding this edge forms a cycle
+        {
+            // remove edge
+            nodes[parent]->removeChild(nodes[child]);
+            nodes[child]->removeParent(nodes[parent]);
+            throw invalid_argument("Adding this edge forms a cycle in the graph.");
+        }
+
     }
-    std::map<std::string, Node*>& Network::getNodes()
+    map<string, Node*>& Network::getNodes()
     {
         return nodes;
     }
-    void Network::fit(const std::vector<std::vector<int>>& dataset, const int smoothing)
+    void Network::fit(const vector<vector<int>>& dataset, const int smoothing)
     {
-        auto jointCounts = [](const std::vector<std::vector<int>>& data, const std::vector<int>& indices, int numStates) {
+        auto jointCounts = [](const vector<vector<int>>& data, const vector<int>& indices, int numStates) {
             int size = indices.size();
-            std::vector<int64_t> sizes(size, numStates);
+            vector<int64_t> sizes(size, numStates);
             torch::Tensor counts = torch::zeros(sizes, torch::kLong);
 
             for (const auto& row : data) {
@@ -41,16 +86,16 @@ namespace bayesnet {
             }
 
             return counts;
-        };
+            };
 
         auto marginalCounts = [](const torch::Tensor& jointCounts) {
             return jointCounts.sum(-1);
-        };
+            };
 
         for (auto& pair : nodes) {
             Node* node = pair.second;
 
-            std::vector<int> indices;
+            vector<int> indices;
             for (const auto& parent : node->getParents()) {
                 indices.push_back(nodes[parent->getName()]->getId());
             }
@@ -67,12 +112,12 @@ namespace bayesnet {
         }
     }
 
-    torch::Tensor& Network::getCPD(const std::string& key)
+    torch::Tensor& Network::getCPD(const string& key)
     {
         return cpds[key];
     }
 
-    void Network::setCPD(const std::string& key, const torch::Tensor& cpt)
+    void Network::setCPD(const string& key, const torch::Tensor& cpt)
     {
         cpds[key] = cpt;
     }