Complete Stratified K Fold

.vscode/settings.json

@@ -98,7 +98,8 @@
         "queue": "cpp",
         "typeindex": "cpp",
         "shared_mutex": "cpp",
-        "*.ipp": "cpp"
+        "*.ipp": "cpp",
+        "cassert": "cpp"
     },
     "cmake.configureOnOpen": false,
     "C_Cpp.default.configurationProvider": "ms-vscode.cmake-tools"

src/Platform/Folding.cc

@@ -34,44 +34,44 @@ pair<vector<int>, vector<int>> KFold::getFold(int nFold)
 StratifiedKFold::StratifiedKFold(int k, const vector<int>& y, int seed) :
     k(k), seed(seed)
 {
-    // n = y.size();
-    // map<int, vector<int>> class_to_indices;
-    // for (int i = 0; i < n; ++i) {
-    //     class_to_indices[y[i]].push_back(i);
-    // }
-    // random_device rd;
-    // default_random_engine random_seed(seed == -1 ? rd() : seed);
-    // for (auto& [cls, indices] : class_to_indices) {
-    //     shuffle(indices.begin(), indices.end(), random_seed);
-    //     int fold_size = n / k;
-    //     for (int i = 0; i < k; ++i) {
-    //         int start = i * fold_size;
-    //         int end = (i == k - 1) ? indices.size() : (i + 1) * fold_size;
-    //         stratified_indices.emplace_back(indices.begin() + start, indices.begin() + end);
-    //     }
-    // }
     n = y.size();
-    stratified_indices.resize(k);
+    stratified_indices = vector<vector<int>>(k);
+    int fold_size = n / k;
+    int remainder = n % k;
+    // Compute class counts and indices
+    auto class_indices = map<int, vector<int>>();
     vector<int> class_counts(*max_element(y.begin(), y.end()) + 1, 0);
     for (auto i = 0; i < n; ++i) {
         class_counts[y[i]]++;
+        class_indices[y[i]].push_back(i);
     }
-    vector<int> class_starts(class_counts.size());
-    partial_sum(class_counts.begin(), class_counts.end() - 1, class_starts.begin() + 1);
-    vector<int> indices(n);
-    for (auto i = 0; i < n; ++i) {
-        int label = y[i];
-        stratified_indices[class_starts[label]] = i;
-        class_starts[label]++;
+    // Shuffle class indices
+    random_device rd;
+    default_random_engine random_seed(seed == -1 ? rd() : seed);
+    for (auto& [cls, indices] : class_indices) {
+        shuffle(indices.begin(), indices.end(), random_seed);
+    }
+    // Assign indices to folds
+    for (auto label = 0; label < class_counts.size(); ++label) {
+        auto num_samples_to_take = class_counts[label] / k;
+        if (num_samples_to_take == 0)
+            continue;
+        auto remainder_samples_to_take = class_counts[label] % k;
+        for (auto fold = 0; fold < k; ++fold) {
+            auto it = next(class_indices[label].begin(), num_samples_to_take);
+            move(class_indices[label].begin(), it, back_inserter(stratified_indices[fold]));  // ##
+            class_indices[label].erase(class_indices[label].begin(), it);
+        }
+        while (remainder_samples_to_take > 0) {
+            int fold = (rand() % static_cast<int>(k));
+            if (stratified_indices[fold].size() == fold_size) {
+                continue;
+            }
+            auto it = next(class_indices[label].begin(), 1);
+            stratified_indices[fold].push_back(*class_indices[label].begin());
+            class_indices[label].erase(class_indices[label].begin(), it);
+            remainder_samples_to_take--;
         }
-    int fold_size = n / k;
-    int remainder = n % k;
-    int start = 0;
-    for (auto i = 0; i < k; ++i) {
-        int fold_length = fold_size + (i < remainder ? 1 : 0);
-        stratified_indices[i].resize(fold_length);
-        copy(indices.begin() + start, indices.begin() + start + fold_length, stratified_indices[i].begin());
-        start += fold_length;
     }
 }
 pair<vector<int>, vector<int>> StratifiedKFold::getFold(int nFold)

src/Platform/Folding.h

@@ -6,8 +6,8 @@ class KFold {
 private:
     int k;
     int n;
-    vector<int> indices;
     int seed;
+    vector<int> indices;
 public:
     KFold(int k, int n, int seed = -1);
     pair<vector<int>, vector<int>> getFold(int nFold);
@@ -16,8 +16,8 @@ class StratifiedKFold {
 private:
     int k;
     int n;
+    int seed;
     vector<vector<int>> stratified_indices;
-    unsigned seed;
 public:
     StratifiedKFold(int k, const vector<int>& y, int seed = -1);
     pair<vector<int>, vector<int>> getFold(int nFold);

src/Platform/testx.cpp

@@ -23,6 +23,7 @@ public:
     C(int a, int c) : A(a), c(c) {}
     int getC() { return c; }
 };
+
 string counts(vector<int> y, vector<int> indices)
 {
     auto result = map<int, int>();
@@ -38,12 +39,12 @@ string counts(vector<int> y, vector<int> indices)
 
 int main()
 {
-	auto y = vector<int>(150);
+    auto y = vector<int>(153);
     fill(y.begin(), y.begin() + 50, 0);
-	fill(y.begin() + 50, y.begin() + 100, 1);
-	fill(y.begin() + 100, y.end(), 2);
+    fill(y.begin() + 50, y.begin() + 103, 1);
+    fill(y.begin() + 103, y.end(), 2);
     //auto fold = KFold(5, 150);
-	auto fold = StratifiedKFold(5, y, 0);
+    auto fold = StratifiedKFold(5, y, -1);
     for (int i = 0; i < 5; ++i) {
         cout << "Fold: " << i << endl;
         auto [train, test] = fold.getFold(i);