Function SplitByClass

catboost/libs/data/objects_grouping.h:257–290

    template <class TClassId>
    TVector<TVector<ui32>> SplitByClass(
        const TObjectsGrouping& objectsGrouping,
        TConstArrayRef<TClassId> objectClasses
    ) {
        CB_ENSURE(objectsGrouping.IsTrivial(), "Stratified split is not supported for data with groups");

        const ui32 objectCount = objectsGrouping.GetObjectCount();

        CheckDataSize(
            objectClasses.size(),
            (size_t)objectCount,
            "objectClasses",
            false,
            "objects size",
            true
        );

        TVector<std::pair<TClassId, ui32>> classWithObject;
        classWithObject.reserve(objectCount);
        for (ui32 i = 0; i < objectCount; ++i) {
            classWithObject.emplace_back(objectClasses[i], i);
        }
        Sort(classWithObject.begin(), classWithObject.end());

        TVector<TVector<ui32>> splittedByClass;
        for (ui32 i = 0; i < classWithObject.size(); ++i) {
            if (i == 0 || classWithObject[i].first != classWithObject[i - 1].first) {
                splittedByClass.emplace_back();
            }
            splittedByClass.back().push_back(classWithObject[i].second);
        }

        return splittedByClass;
    }

    template <class TClassId>
    void StratifiedTrainTestSplit(