1 files changed, 11 insertions, 3 deletions

diff --git a/split.cpp b/split.cpp
index be80545..e774ea9 100644
--- a/split.cpp
+++ b/split.cpp
@@ -43,6 +43,7 @@ int main(int argc, char **argv)
     desc.add_options()
             ("help,h", "produce help message")
             ("query,Q", "Query person count")
+            ("shuffle,s","shuffle input")
             ("count,c", po::value<int>(), "take number")
             ("take,a", po::value<string>(), "set output rid file 1(taken)")
             ("left,b", po::value<string>(), "set output rid file 2(left)")
@@ -68,9 +69,16 @@ int main(int argc, char **argv)
     vector<DataEntry*> a;
     vector<DataEntry*> b;
     dp.open();
-    dp.take(vm["count"].as<int>(),a,b);
-    outputRid(a,dp.getfSize(),vm["take"].as<string>());
-    outputRid(b,dp.getfSize(),vm["left"].as<string>());
+    if (vm.count("shuffle")) {
+        dp.getRidVector(a);
+        dp.shuffle(a);
+        outputRid(a, dp.getfSize(), vm["take"].as<string>());
+    }
+    else {
+        dp.take(vm["count"].as<int>(), a, b);
+        outputRid(a, dp.getfSize(), vm["take"].as<string>());
+        outputRid(b, dp.getfSize(), vm["left"].as<string>());
+    }
     dp.close();
     return 0;
 }
\ No newline at end of file