Quick Select

lib/cpalgo/util/README.md

@@ -50,4 +50,14 @@ Arithmetic in Modulo
 
 ### References in Japanese
 - [「1000000007 で割ったあまり」の求め方を総特集！ 〜 逆元から離散対数まで 〜 - Qiita](https://qiita.com/drken/items/3b4fdf0a78e7a138cd9a#8-modint)
-- [modint 構造体を使ってみませんか？ (C++) - noshi91のメモ](https://noshi91.hatenablog.com/entry/2019/03/31/174006)
+- [modint 構造体を使ってみませんか？ (C++) - noshi91のメモ](https://noshi91.hatenablog.com/entry/2019/03/31/174006)
+
+
+## Quick Select
+
+Selects the k-th smallest element of an unsorted array in linear time.
+
+[quick_select | C++ code](quick_select.hpp)
+
+### References in English
+- [Median of medians - Wikipedia](https://en.wikipedia.org/wiki/Median_of_medians)

lib/cpalgo/util/quick_select.hpp

@@ -0,0 +1,93 @@
+#pragma once
+
+#include <algorithm>
+#include <stdexcept>
+#include <vector>
+
+
+// Selects the k-th smallest element of an unsorted array.
+//
+// Complexity:
+//   Time: O(N)
+//   Space: O(logN) ~ O(N)
+//
+template <class T, class Comparator = std::less<T>>
+struct QuickSelect {
+    Comparator less = Comparator();
+    // Selects the k-th smallest element of the given unsorted array.
+    // Returns the value of the k-th smallest element.
+    // k = [0,N)
+    // Time: O(N)
+    // Space: O(N)
+    T operator()(std::vector<T> const& array, size_t k) const {
+        if (k < 0 || k >= array.size()) throw std::out_of_range("k");
+        std::vector<T> buffer(array.begin(), array.end());
+        auto index = index_of(buffer, k);
+        return buffer[index];
+    }
+    // Selects the k-th smallest element of the given unsorted array.
+    // Returns the index of the k-th smallest element.
+    // k = [0,N)
+    // Time: O(N)
+    // Space: O(logN)
+    size_t index_of(std::vector<T>& array, size_t k) const {
+        if (k < 0 || k >= array.size()) throw std::out_of_range("k");
+        return select(array, 0, array.size(), k);
+    }
+private:
+    size_t select(std::vector<T>& array, size_t begin, size_t end, size_t k) const {
+        while (true) {
+            if (end - begin == 1) {
+                return begin;
+            }
+            auto pivot_index = pivot_of(array, begin, end);
+            pivot_index = partition(array, begin, end, k, array[pivot_index]);
+            std::cout << pivot_index << " " << k << std::endl;
+            if (pivot_index == k) {
+                return k;
+            } else if (pivot_index < k) {
+                begin = pivot_index + 1;
+            } else {
+                end = pivot_index;
+            }
+        }
+    }
+    size_t pivot_of(std::vector<T>& array, size_t begin, size_t end) const {
+        if (end - begin <= 5) {
+            return partition5(array, begin, end);
+        }
+        size_t n = 0;
+        for (size_t i = begin; i < end; i += 5) {
+            auto median = partition5(array, i, std::min(i+5,end));
+            std::swap(array[median], array[begin + (i-begin)/5]);
+            ++n;
+        }
+        auto k = begin + (n-1)/2;
+        return select(array, begin, begin + n, k);
+    }
+    inline size_t partition(std::vector<T>& array, size_t begin, size_t end, size_t k, T pivot_value) const {
+        auto begin_equal = std::partition(array.begin() + begin, array.begin() + end, [&](auto value) {
+            return less(value, pivot_value);
+        });
+        auto end_equal = std::partition(begin_equal, array.begin() + end, [&](auto value) {
+            return !less(value, pivot_value) && !less(pivot_value, value);
+        });
+        size_t begin_equal_range = distance(array.begin(), begin_equal);
+        size_t end_equal_range = distance(array.begin(), end_equal);
+        if (k < begin_equal_range) {
+            return begin_equal_range;
+        } else if (k >= end_equal_range) {
+            return end_equal_range-1;
+        } else {
+            return k;
+        }
+    }
+    inline size_t partition5(std::vector<T>& array, size_t begin, size_t end) const {
+        for (size_t i = begin + 1; i < end; ++i) {
+            for (size_t j = i; j > begin && less(array[j], array[j-1]); --j) {
+                std::swap(array[j-1], array[j]);
+            }
+        }
+        return begin + (end - begin - 1) / 2;
+    }
+};

lib/cpalgo/util/tests/quick_select_test.cpp

@@ -0,0 +1,79 @@
+#include <bits/stdc++.h>
+#include "gtest/gtest.h"
+#include "cpalgo/util/quick_select.hpp"
+
+using namespace std;
+
+TEST(QuickSelectTest, ShouldSelectInSmallCases) {
+
+    auto select = QuickSelect<int>();
+    vector<int> n1 = { 1 };
+    EXPECT_EQ( 1, select(n1, 0));
+
+    vector<int> n12 = { 1, 2 };
+    EXPECT_EQ( 1, select(n12, 0));
+    EXPECT_EQ( 2, select(n12, 1));
+
+    vector<int> n11 = { 1, 1 };
+    EXPECT_EQ( 1, select(n11, 0));
+    EXPECT_EQ( 1, select(n11, 1));
+
+    vector<int> n112 = { 1, 1, 2 };
+    EXPECT_EQ( 1, select(n112, 0));
+    EXPECT_EQ( 1, select(n112, 1));
+    EXPECT_EQ( 2, select(n112, 2));
+
+}
+
+TEST(QuickSelectTest, ShouldSelect) {
+
+    auto select = QuickSelect<int>();
+    vector<int> values = { 1, 3, 2, 6, 7, 4, 5, 10, 8, 9 };
+    EXPECT_EQ( 1, select(values, 0));
+    EXPECT_EQ( 2, select(values, 1));
+    EXPECT_EQ( 3, select(values, 2));
+    EXPECT_EQ( 4, select(values, 3));
+    EXPECT_EQ( 5, select(values, 4));
+    EXPECT_EQ( 6, select(values, 5));
+    EXPECT_EQ( 7, select(values, 6));
+    EXPECT_EQ( 8, select(values, 7));
+    EXPECT_EQ( 9, select(values, 8));
+    EXPECT_EQ(10, select(values, 9));
+
+}
+
+TEST(QuickSelectTest, ShouldSelectUsingCustomComparator) {
+
+    auto select = QuickSelect<int, greater<int>>();
+    vector<int> values = { 1, 3, 2, 6, 7, 4, 5, 10, 8, 9 };
+    EXPECT_EQ(10, select(values, 0));
+    EXPECT_EQ( 9, select(values, 1));
+    EXPECT_EQ( 8, select(values, 2));
+    EXPECT_EQ( 7, select(values, 3));
+    EXPECT_EQ( 6, select(values, 4));
+    EXPECT_EQ( 5, select(values, 5));
+    EXPECT_EQ( 4, select(values, 6));
+    EXPECT_EQ( 3, select(values, 7));
+    EXPECT_EQ( 2, select(values, 8));
+    EXPECT_EQ( 1, select(values, 9));
+
+}
+
+TEST(QuickSelectTest, ShouldSelectInRandomizedCase) {
+
+    size_t const N = 1000;
+    vector<int> values(N);
+    random_device seed_generator;
+    mt19937 random(seed_generator());
+    for (auto &value : values) {
+        value = abs((int)(random() % 100));
+    }
+    auto sorted_values = values;
+    sort(sorted_values.begin(), sorted_values.end());
+
+    auto select = QuickSelect<int>();
+    for (size_t i = 0; i < N; ++i) {
+        EXPECT_EQ(sorted_values[i], select(values, i));
+    }
+
+}