From 3750662f2c55275fec3adbd43f606ce4c53d13e4 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Ricardo=20Montan=CC=83ana?= <rmontanana@gmail.com>
Date: Sun, 9 Jul 2023 16:25:24 +0200
Subject: [PATCH] Add getStates method

---
 src/Network.cc | 8 ++++++++
 src/Network.h  | 1 +
 2 files changed, 9 insertions(+)

diff --git a/src/Network.cc b/src/Network.cc
index 391b8af..4c039b4 100644
--- a/src/Network.cc
+++ b/src/Network.cc
@@ -38,6 +38,14 @@ namespace bayesnet {
     {
         return classNumStates;
     }
+    int Network::getStates()
+    {
+        int result = 0;
+        for (auto node : nodes) {
+            result += node.second->getNumStates();
+        }
+        return result;
+    }
     string Network::getClassName()
     {
         return className;
diff --git a/src/Network.h b/src/Network.h
index 78ce3ab..7db04f9 100644
--- a/src/Network.h
+++ b/src/Network.h
@@ -30,6 +30,7 @@ namespace bayesnet {
         void addEdge(const string, const string);
         map<string, Node*>& getNodes();
         vector<string> getFeatures();
+        int getStates();
         int getClassNumStates();
         string getClassName();
         void fit(const vector<vector<int>>&, const vector<int>&, const vector<string>&, const string&);