Function StratifiedTrainTestSplit

catboost/libs/data/objects_grouping.h:293–325

    template <class TClassId>
    void StratifiedTrainTestSplit(
        const TObjectsGrouping& objectsGrouping,
        TConstArrayRef<TClassId> objectClasses,
        double trainPart,
        TArraySubsetIndexing<ui32>* trainIndices,
        TArraySubsetIndexing<ui32>* testIndices
    ) {
        TVector<TVector<ui32>> splittedByClass = SplitByClass(objectsGrouping, objectClasses);
        ui32 minLen = GetClassSplitMinLen(objectsGrouping.GetObjectCount(), splittedByClass);
        if (minLen < 2) {
            CATBOOST_WARNING_LOG << " Warning: The least populated class in y has only "
                << minLen << " members, which is too few.";
        }
        TVector<ui32> resultTrainIndices;
        TVector<ui32> resultTestIndices;
        for (const auto& part : splittedByClass) {
            for (ui32 idx = 0; idx < part.size() * trainPart; ++idx) {
                resultTrainIndices.push_back(part[idx]);
            }
            for (ui32 idx = part.size() * trainPart; idx < part.size(); ++idx) {
                resultTestIndices.push_back(part[idx]);
            }
        }

        CB_ENSURE(!resultTrainIndices.empty(), "Not enough objects for splitting into train and test subsets");
        CB_ENSURE(!resultTestIndices.empty(), "Not enough objects for splitting into train and test subsets");

        Sort(resultTrainIndices.begin(), resultTrainIndices.end());
        *trainIndices = TArraySubsetIndexing<ui32>(std::move(resultTrainIndices));

        Sort(resultTestIndices.begin(), resultTestIndices.end());
        *testIndices = TArraySubsetIndexing<ui32>(std::move(resultTestIndices));
    }

    template <class TClassId>
    TVector<TArraySubsetIndexing<ui32>> StratifiedSplitToFolds(