Splitting std::vector based on some criteria

Question

I have a vector which contain some data. I want to split it into const number of vectors depending on some criteria. For example:

using Point=std::pair<int,int>;
std::array<std::vector<Point>,4> split_to_4(const std::vector<Point>& data,std::function<size_t(Point)> criteria);
int main(){
    std::vector<Point> data;
    //fill data
    auto results=split_to_4(data,[](const Point& p){
        if(cond1) return 0;
        if(cond2) return 1;
        if(cond3) return 2;
        return 3;
    });
}

What is the best way to implement split_to_4? My current attempt is:

std::array<std::vector<Point>,4> split_to_4(const std::vector<Point>& data,std::function<size_t(Point)> criteria){
    std::array<std::vector<Point>,4> result;
    for (const auto& p : data){
        areas_regions[criteria(p)].emplace_back(p);
    }
    return result;
}

Any better.. More std way to do it?

By Better, I mean: more readable... depend on iterator... depend on some std functions...

Answer 1

You can do this in place with multiple calls to std::partition:

// Returns iterators to the three partition points in the range
template<class ForwardIt, class Which>
auto split4(ForwardIt first, ForwardIt last, Which which) {
    std::array<ForwardIt, 3> ret;
    ret[0] = std::partition(first, last, 
                [&](const auto &v){return which(v) == 0;});
    ret[1] = std::partition(ret[0], last, 
                [&](const auto &v){return which(v) == 1;});
    ret[2] = std::partition(ret[1], last,
                [&](const auto &v){return which(v) == 2;});
    return ret;
}

Of course you can also pass and use the conditions directly instead of proxiing through some which function if you so desire.

One could also trivially rewrite this with a loop to generalize it to splitN if necessary. (Watch out though, the complexity of this approach is O(N * n) for a range with n elements. This will probably be unreasonably slow for big N. On the other hand, we get swaps instead of copies, which may help if copying is expensive (compared to calling which). If performance is critical, measure.)

If you need the relative order of elements in each group preserved, std::stable_partition is your friend.

---

Just noticed the C++11 tag: The above code is C++14. For C++11 compatibility, simply change the autos I used to the explicit types, i.e. use std::array<ForwardIt, 3> as the return type and const std::iterator_traits<ForwardIt>::value_type& for the lambdas.

I'll leave that as is for brevity, this last paragraph completes the answer for the pre-C++14 folks.

Answer 2

update:

probably the most STL-like way:

Features:

1. Iterator-based so choice of source and destination containers is left to the caller

2. Source iterators may be move-iterators if move-partitioning is required, or leave as normal iterators to make a copy

3. Linear time complexity

4. Stable ordering of results (ref std::stable_partition)

-

#include <array>
#include <vector>
#include <utility>
#include <cassert>

using Point=std::pair<int,int>;

// example split function - could be a function object
extern std::size_t which_bucket(const Point&);


template<class Iter, class OutIter, class Which>
  auto split_n(Iter first, Iter last, 
               OutIter outfirst, std::size_t N, 
              Which&& which)
{
  while (first != last) {
    auto index = which(*first);
    assert (index < N);
    std::next(outfirst, index) -> push_back(*first);
    ++ first;
  }
}

template<class Iter, class OutIter, class Which>
  auto split_to(Iter first, Iter last, 
               OutIter outfirst, OutIter outlast, 
              Which&& which)
{
  return split_n(first, last, outfirst, 
                  std::distance(outfirst, outlast),
                  std::forward<Which>(which));
}


int main(){
  std::vector<Point> source;
  std::array<std::vector<Point>, 4> dest { }; 

  split_n(source.begin(), source.end(), 
          dest.begin(), dest.size(), 
          which_bucket);

  // or

  split_to(source.begin(), source.end(),
           dest.begin(), dest.end(),
           which_bucket);

  // or with move request:

  split_to(std::make_move_iterator(source.begin()), 
           std::make_move_iterator(source.end()),
           dest.begin(), dest.end(),
           which_bucket);
}

another way

#include <array>
#include <vector>
#include <utility>

using Point=std::pair<int,int>;

// example split function - could be a function object
extern std::size_t which_bucket(const Point&);


template<class Iter, class Which>
  auto split4(Iter first, Iter last, Which&& which)
{
  std::array<std::vector<Point>, 4> result {};
  while (first != last) {
    result[which(*first)].push_back(*first);
    ++first;
  }
  return result;
}

int main(){
  std::vector<Point> data;

  auto results = split4(data.begin(), data.end(), which_bucket);
}

Here's another way which honours any custom allocator in the vector:

#include <array>
#include <vector>
#include <utility>

using Point=std::pair<int,int>;

// example split function - could be a function object
extern std::size_t which_bucket(const Point&);


template<class T, class A, class Which>
  auto split4(const std::vector<T,A>& v, 
              Which&& which)
{
  using vec_type = std::vector<T,A>;
  std::array<std::vector<T,A>, 4> result { 
    vec_type(v.get_allocator()),
    vec_type(v.get_allocator()),
    vec_type(v.get_allocator()),
    vec_type(v.get_allocator())
  };
  for (auto& p : v) {
    result[which(p)].push_back(p);
  }
  return result;
}

int main(){
  std::vector<Point> data;

  auto results = split4(data, which_bucket);
}