[REFACT] Refactored entropy calculator allowing to supply histogram

hasherezade · Jul 6, 2023 · e2b1525 · e2b1525
1 parent c9f87f0
commit e2b1525
Showing 1 changed file with 26 additions and 13 deletions.
diff --git a/stats/entropy.h b/stats/entropy.h
@@ -5,26 +5,39 @@ namespace pesieve {
 
     namespace util {
 
-        // from: https://stackoverflow.com/questions/20965960/shannon-entropy
-        template <typename T> static double ShannonEntropy(T data[], size_t elements)
+        template <typename T> size_t generateHistogram(IN T buffer[], IN size_t bufferSize, OUT std::map<T, size_t> &counts)
         {
-            double entropy = 0;
-            std::map<T, size_t> counts;
-            typename std::map<T, size_t>::iterator it;
-            //
-            for (size_t dataIndex = 0; dataIndex < elements; ++dataIndex) {
-                const T val = data[dataIndex];
+            if (!buffer || !bufferSize) return 0;
+
+            for (size_t i = 0; i < bufferSize; ++i) {
+                const T val = buffer[i];
                 counts[val]++;
             }
-            if (!elements) return 0;
-            //
-            for (it = counts.begin(); it != counts.end(); ++it) {
-                double p_x = (double)it->second / elements;
-                if (p_x > 0) entropy -= p_x * log(p_x) / log((double)2);
+            return counts.size();
+        }
+
+        // Shannon's Entropy calculation based on: https://stackoverflow.com/questions/20965960/shannon-entropy
+        template <typename T>
+        double calcShannonEntropy(std::map<T, size_t>& histogram, size_t totalSize)
+        {
+            if (!totalSize) return 0;
+            double entropy = 0;
+            for (auto it = histogram.begin(); it != histogram.end(); ++it) {
+                double p_x = (double)it->second / totalSize;
+                if (p_x > 0) entropy -= p_x * log(p_x) / log(2);
             }
             return entropy;
         }
 
+        template <typename T> static double ShannonEntropy(T buffer[], size_t bufferSize)
+        {
+            std::map<T, size_t> counts;
+            if (!generateHistogram<T>(buffer, bufferSize, counts)) {
+                return 0;
+            }
+            return calcShannonEntropy<T>(counts, bufferSize);
+        }
+
     }; // namespace util
 
 }; //namespace pesieve