%PDF-
%PDF-
Mini Shell
Mini Shell
#ifndef PYTHONIC_NUMPY_ARRAYSPLIT_HPP
#define PYTHONIC_NUMPY_ARRAYSPLIT_HPP
#include "pythonic/include/numpy/array_split.hpp"
#include "pythonic/utils/functor.hpp"
#include "pythonic/types/ndarray.hpp"
PYTHONIC_NS_BEGIN
namespace numpy
{
template <class T, class pS>
types::list<types::ndarray<T, types::array<long, std::tuple_size<pS>::value>>>
array_split(types::ndarray<T, pS> const &a, long nb_split)
{
long sz = std::distance(a.begin(), a.end());
long n = (sz + nb_split - 1) / nb_split;
long end = n * nb_split;
long nb_full_split = nb_split;
if (end != sz)
nb_full_split -= (end - sz);
types::list<types::ndarray<
T, types::array<long, std::tuple_size<pS>::value>>> out(nb_split);
long index = 0;
for (long i = 0; i < nb_full_split; ++i, index += n)
out[i] = a[types::contiguous_slice(index, index + n)];
for (long i = nb_full_split; i < nb_split; ++i, index += (n - 1))
out[i] = a[types::contiguous_slice(index, index + n - 1)];
return out;
}
template <class T, class pS, class I>
typename std::enable_if<
types::is_iterable<I>::value,
types::list<types::ndarray<
T, types::array<long, std::tuple_size<pS>::value>>>>::type
array_split(types::ndarray<T, pS> const &a, I const &split_mask)
{
long sz = std::distance(a.begin(), a.end());
types::list<
types::ndarray<T, types::array<long, std::tuple_size<pS>::value>>>
out(1 + split_mask.flat_size());
long index = 0;
auto inserter = out.begin();
for (auto next_index : split_mask) {
*inserter++ = a[types::contiguous_slice(index, next_index)];
index = next_index;
}
*inserter = a[types::contiguous_slice(index, sz)];
return out;
}
NUMPY_EXPR_TO_NDARRAY0_IMPL(array_split);
}
PYTHONIC_NS_END
#endif
Zerion Mini Shell 1.0