arrow/util/sort.h

0001 // Licensed to the Apache Software Foundation (ASF) under one
0002 // or more contributor license agreements.  See the NOTICE file
0003 // distributed with this work for additional information
0004 // regarding copyright ownership.  The ASF licenses this file
0005 // to you under the Apache License, Version 2.0 (the
0006 // "License"); you may not use this file except in compliance
0007 // with the License.  You may obtain a copy of the License at
0008 //
0009 //   http://www.apache.org/licenses/LICENSE-2.0
0010 //
0011 // Unless required by applicable law or agreed to in writing,
0012 // software distributed under the License is distributed on an
0013 // "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
0014 // KIND, either express or implied.  See the License for the
0015 // specific language governing permissions and limitations
0016 // under the License.
0017
0018 #pragma once
0019
0020 #include <algorithm>
0021 #include <cstdint>
0022 #include <functional>
0023 #include <numeric>
0024 #include <utility>
0025 #include <vector>
0026
0027 namespace arrow {
0028 namespace internal {
0029
0030 template <typename T, typename Cmp = std::less<T>>
0031 std::vector<int64_t> ArgSort(const std::vector<T>& values, Cmp&& cmp = {}) {
0032   std::vector<int64_t> indices(values.size());
0033   std::iota(indices.begin(), indices.end(), 0);
0034   std::sort(indices.begin(), indices.end(),
0035             [&](int64_t i, int64_t j) -> bool { return cmp(values[i], values[j]); });
0036   return indices;
0037 }
0038
0039 template <typename T>
0040 size_t Permute(const std::vector<int64_t>& indices, std::vector<T>* values) {
0041   if (indices.size() <= 1) {
0042     return indices.size();
0043   }
0044
0045   // mask indicating which of values are in the correct location
0046   std::vector<bool> sorted(indices.size(), false);
0047
0048   size_t cycle_count = 0;
0049
0050   for (auto cycle_start = sorted.begin(); cycle_start != sorted.end();
0051        cycle_start = std::find(cycle_start, sorted.end(), false)) {
0052     ++cycle_count;
0053
0054     // position in which an element belongs WRT sort
0055     auto sort_into = static_cast<int64_t>(cycle_start - sorted.begin());
0056
0057     if (indices[sort_into] == sort_into) {
0058       // trivial cycle
0059       sorted[sort_into] = true;
0060       continue;
0061     }
0062
0063     // resolve this cycle
0064     const auto end = sort_into;
0065     for (int64_t take_from = indices[sort_into]; take_from != end;
0066          take_from = indices[sort_into]) {
0067       std::swap(values->at(sort_into), values->at(take_from));
0068       sorted[sort_into] = true;
0069       sort_into = take_from;
0070     }
0071     sorted[sort_into] = true;
0072   }
0073
0074   return cycle_count;
0075 }
0076
0077 }  // namespace internal
0078 }  // namespace arrow