select.hpp

Go to the documentation of this file.

#pragma once
 
#include <cstdlib>
#include <initializer_list>
#include <limits>
#include <vector>
 
#include "NumCpp/Core/Internal/Error.hpp"
#include "NumCpp/Core/Internal/StlAlgorithms.hpp"
#include "NumCpp/Core/Shape.hpp"
#include "NumCpp/NdArray.hpp"
 
namespace nc
{
    //============================================================================
    // Method Description:
    template<typename dtype>
    NdArray<dtype> select(const std::vector<const NdArray<bool>*>&  condVec,
                          const std::vector<const NdArray<dtype>*>& choiceVec,
                          dtype                                     defaultValue = dtype{ 0 })
    {
        if (choiceVec.size() != condVec.size())
        {
            THROW_INVALID_ARGUMENT_ERROR("condVec and choiceVec need to be the same size");
        }
 
        if (choiceVec.size() == 0)
        {
            THROW_INVALID_ARGUMENT_ERROR("choiceVec is size 0");
        }
 
        auto theShape = condVec.front()->shape();
        for (const auto cond : condVec)
        {
            const auto& theCond = *cond;
            if (theCond.shape() != theShape)
            {
                THROW_INVALID_ARGUMENT_ERROR("all NdArrays of the condVec must be the same shape");
            }
        }
 
        for (const auto choice : choiceVec)
        {
            const auto& theChoice = *choice;
            if (theChoice.shape() != theShape)
            {
                THROW_INVALID_ARGUMENT_ERROR(
                    "all NdArrays of the choiceVec must be the same shape, and the same as condVec");
            }
        }
 
        using size_type           = typename NdArray<dtype>::size_type;
        constexpr auto nullChoice = std::numeric_limits<size_type>::max();
 
        NdArray<size_type> choiceIndices(theShape);
        choiceIndices.fill(nullChoice);
        for (size_type condIdx = 0; condIdx < condVec.size(); ++condIdx)
        {
            const auto& theCond = *condVec[condIdx];
            for (size_type i = 0; i < theCond.size(); ++i)
            {
                if (theCond[i] && choiceIndices[i] == nullChoice)
                {
                    choiceIndices[i] = condIdx;
                }
            }
        }
 
        NdArray<dtype> result(theShape);
        result.fill(defaultValue);
        for (size_type i = 0; i < choiceIndices.size(); ++i)
        {
            const auto choiceIndex = choiceIndices[i];
            if (choiceIndex != nullChoice)
            {
                const auto& theChoice = *choiceVec[choiceIndex];
                result[i]             = theChoice[i];
            }
        }
 
        return result;
    }
 
    //============================================================================
    // Method Description:
    template<typename dtype>
    NdArray<dtype> select(const std::vector<NdArray<bool>>&  condList,
                          const std::vector<NdArray<dtype>>& choiceList,
                          dtype                              defaultValue = dtype{ 0 })
    {
        std::vector<const NdArray<bool>*> condVec(condList.size());
        stl_algorithms::transform(condList.begin(),
                                  condList.end(),
                                  condVec.begin(),
                                  [](auto& cond) noexcept -> const NdArray<bool>* { return &cond; });
 
        std::vector<const NdArray<dtype>*> choiceVec(choiceList.size());
        stl_algorithms::transform(choiceList.begin(),
                                  choiceList.end(),
                                  choiceVec.begin(),
                                  [](auto& choice) noexcept -> const NdArray<dtype>* { return &choice; });
 
        return select(condVec, choiceVec, defaultValue);
    }
} // namespace nc