SelectRandomWeighted.h

Go to the documentation of this file.

// Copyright Epic Games, Inc. All Rights Reserved.
 
#pragma once
 
#include "Algo/Accumulate.h"
#include "Algo/Impl/RangePointerType.h"
#include "Math/RandomStream.h"
#include "Templates/UnrealTemplate.h"
#include "Traits/ElementType.h"
#include "GenericPlatform/GenericPlatformMath.h"
 
namespace AlgoImpl
{
    template <typename RangeType, typename ProjectionType>
    [[nodiscard]] typename TRangePointerType<typename TRemoveReference<RangeType>::Type>::Type SelectRandomWeightedBy(RangeType&& Range, ProjectionType Proj)
    {
        using WeightType = std::decay_t<TInvokeResult_T<ProjectionType, TElementType_T<RangeType>>>;
 
        const auto SumOfAllDesires = Algo::Accumulate(Range, (WeightType)0, [&Proj](auto Acc, auto&& Elem)
        {
            const auto Weight = Invoke(Proj, Elem);
            // Negative values are invalid and should be ignored
            return Acc + (Weight < (WeightType)0 ? (WeightType)0 : Weight);
        });
        
        auto RandomWeightedAvg = (WeightType)(FMath::FRand() * SumOfAllDesires);
 
        for (auto&& Elem : Forward<RangeType>(Range))
        {
            const auto Weight = Invoke(Proj, Elem);
 
            // Negative- or zero-weighted elements are never chosen, and are not subtracted from the total since they are not added above.
            if (Weight <= (WeightType)0)
            {
                continue;
            }
 
            if (RandomWeightedAvg < Weight)
            {
                return &Elem;
            }
 
            RandomWeightedAvg -= Weight;
        }
 
        return nullptr;
    }
}
 
namespace Algo
{
    template <typename RangeType, typename ProjectionType>
    UE_REWRITE auto SelectRandomWeightedBy(RangeType&& Range, ProjectionType Proj)
        -> decltype(AlgoImpl::SelectRandomWeightedBy(Forward<RangeType>(Range), MoveTemp(Proj)))
    {
        return AlgoImpl::SelectRandomWeightedBy(Forward<RangeType>(Range), MoveTemp(Proj));
    }
}