From 76fe709d67630ebbefd54f019f0e5eb9dba13e62 Mon Sep 17 00:00:00 2001
From: Anush V Kini <anushkini@gmail.com>
Date: Sun, 1 Nov 2020 20:25:28 +0530
Subject: [PATCH] Changed unordered map implementation to uvec implementation

---
 src/mlpack/core/data/split_data.hpp | 24 +++++++++++++-----------
 1 file changed, 13 insertions(+), 11 deletions(-)

diff --git a/src/mlpack/core/data/split_data.hpp b/src/mlpack/core/data/split_data.hpp
index 66b7cce5cf..2ecb5d720f 100644
--- a/src/mlpack/core/data/split_data.hpp
+++ b/src/mlpack/core/data/split_data.hpp
@@ -14,7 +14,6 @@
 #define MLPACK_CORE_DATA_SPLIT_DATA_HPP
 
 #include <mlpack/prereqs.hpp>
-#include <unordered_map>
 
 namespace mlpack {
 namespace data {
@@ -71,14 +70,14 @@ void StratifiedSplit(const arma::Mat<T>& input,
    * The number of 1 labels in our test set = floor(11 * 0.2) = 2.
    *
    * In our first pass over the dataset,
-   * we visit each label and keep the count of each label in our unordered map.
+   * We visit each label and keep count of each label in our 'labelMap' uvec.
    *
    * We then take a second pass over the dataset.
-   * We now maintain an additional unordered map to hold the label counts in
-   * our test set.
+   * We now maintain an additional uvec 'testLabelMap' to hold the label counts
+   * of our test set.
    *
-   * In this pass, when we encounter a label we check the test set map for
-   * the count of this label in the test set.
+   * In this pass, when we encounter a label we check the 'testLabelMap' uvec
+   * for the count of this label in the test set.
    * If this count is less than the required number of labels in the test set,
    * we add the data to the test set and increment the label count in the map.
    * If this count is equal to or more than the required count in the test set,
@@ -97,9 +96,12 @@ void StratifiedSplit(const arma::Mat<T>& input,
   size_t testIdx = 0;
   size_t trainSize = 0;
   size_t testSize = 0;
+  arma::uvec labelMap;
+  arma::uvec testLabelMap;
+  U maxLabel = inputLabel.max();
 
-  std::unordered_map<U, size_t> labelMap;
-  std::unordered_map<U, size_t> testLabelMap;
+  labelMap.zeros(maxLabel+1);
+  testLabelMap.zeros(maxLabel+1);
 
   arma::uvec order =
       arma::linspace<arma::uvec>(0, input.n_cols - 1, input.n_cols);
@@ -115,10 +117,10 @@ void StratifiedSplit(const arma::Mat<T>& input,
     labelMap[label] += 1;
   }
 
-  for (std::pair<U, size_t> countPair : labelMap)
+  for (arma::uword labelCount : labelMap)
   {
-      testSize += floor(countPair.second * testRatio);
-      trainSize += countPair.second - floor(countPair.second * testRatio);
+      testSize += floor(labelCount * testRatio);
+      trainSize += labelCount - floor(labelCount * testRatio);
   }
 
   trainData.set_size(input.n_rows, trainSize);