Repository: opencv/opencv
File: modules/dnn/src/layers/topk2_layer.cpp (C++)

Code before:
        outputs[0] = inputs.front();
        outputs[1] = CV_64S;
    }

private:
    template<class Comparator, typename T>
    void FindTopK(const Mat &input, Mat &output_vals, Mat &output_idxs, int normalized_axis)
    {
        const auto inShape = shape(input);
        size_t outer = std::accumulate(inShape.begin(), inShape.begin() + normalized_axis, 1, std::multiplies<int>());
        size_t inner = std::accumulate(inShape.begin() + normalized_axis + 1, inShape.end(), 1, std::multiplies<int>());
        int dimAxis = inShape[normalized_axis];

        auto worker = [&](const Range &r) {
            const T* inPtr = input.ptr<T>();
            T* valPtr = output_vals.ptr<T>();
            int64_t* idxPtr = output_idxs.ptr<int64_t>();
            Comparator cmp(inPtr, inner);
            AutoBuffer<size_t> indices(dimAxis);
            size_t* idxBuf = indices.data();
            for (size_t b = r.start; b < r.end; ++b) {
                for (size_t j = 0; j < inner; ++j) {
                    size_t offset = b * dimAxis * inner + j;
                    cmp.addOffset(offset);
                    std::iota(idxBuf, idxBuf + dimAxis, 0);
                    std::stable_sort(idxBuf, idxBuf + dimAxis, cmp);
                    // By the time we reach here, 'K' has been resolved. Use it directly.
                    int currentK = K;
                    for (int i = 0; i < currentK; ++i) {
                        size_t src = idxBuf[i];
                        valPtr[b * currentK * inner + i * inner + j] = inPtr[offset + src * inner];
                        idxPtr[b * currentK * inner + i * inner + j] = static_cast<int64_t>(src);
                    }
                    cmp.minusOffset(offset);
                }
            }
        };
        parallel_for_(Range(0, static_cast<int>(outer)), worker);
    }

public:
    void forward(InputArrayOfArrays inputs_arr,
                 OutputArrayOfArrays outputs_arr,
                 OutputArrayOfArrays internals_arr) CV_OVERRIDE
    {
        CV_TRACE_FUNCTION();
        std::vector<Mat> inputs, outputs;
        inputs_arr.getMatVector(inputs);
        outputs_arr.getMatVector(outputs);

        const auto &input = inputs.front();

Review comment (suggestion):
you don't need std::stable_sort, because you take into account index order inside the comparator class. And you don't need std::sort, because you don't need completely sorted array. Use std::partial_sort.

Also, sorting is done on array that is not continuous in memory. Please, copy elements to a temporary vector (make another AutoBuffer nearby indices) and sort them there. Actually, I recommend you to use std::vector<> of pairs of elements:

```
std::vector<std::pair<T /* or WT, see comment below */, int64_t> > sorted;
for (b...) {
   for (j...) {
        for (int64_t i = 0; i < dimAxis; i++)
            sorted[i] = std::make_pair(data[i*... + ...], i);
        std::partial_sort(sorted.begin(), sorted.begin() + K, sorted.end(), cmp);
        // store K selected values and the corresponding indices from sorted to the target tensor
        ...
   }
}
```


Code after:
                    size_t offset = b * dimAxis * inner + j;
                    for (uint32_t u = 0; u < (uint32_t)dimAxis; ++u) {
                        sortbuf[u].first  = u;
                        sortbuf[u].second = WT(inPtr[offset + u * inner]);
                    }
                    if (largest){
                        ComparatorGreater<uint32_t,WT> cmp;
                        std::partial_sort(sortbuf.begin(), sortbuf.begin() + kVal, sortbuf.end(), cmp);
                    }
                    else{
                        ComparatorLess<uint32_t,WT> cmp;
                        std::partial_sort(sortbuf.begin(), sortbuf.begin() + kVal, sortbuf.end(), cmp);
                    }
                    for (int i = 0; i < kVal; ++i) {
                        auto &p = sortbuf[i];
                        valPtr[b * kVal * inner + i * inner + j] = T(p.second);
                        idxPtr[b * kVal * inner + i * inner + j] = p.first;
                    }
                }
            }
        };
        parallel_for_(Range(0, static_cast<int>(outer)), worker);
    }

public:
    void forward(InputArrayOfArrays inputs_arr,
                 OutputArrayOfArrays outputs_arr,
                 OutputArrayOfArrays internals_arr) CV_OVERRIDE
    {
        CV_TRACE_FUNCTION();
        std::vector<Mat> inputs, outputs;
        inputs_arr.getMatVector(inputs);
        outputs_arr.getMatVector(outputs);

        const auto &input = inputs.front();
        Mat output_value, output_index;

        // Normalize axis to handle negative values
        int normalized_axis = normalize_axis(axis, input.dims);

        int kVal = K;
        if (dynamicK) {
            CV_Assert(inputs.size() == 2);
            CV_Assert(inputs[1].type() == CV_64S);
            CV_Assert(inputs[1].total() == 1);
            int64_t kTemp = inputs[1].at<int64_t>(0);
            CV_CheckGT(kTemp, 0, "TopK2: dynamic K must be > 0");
            CV_CheckLT(kTemp, input.size[normalized_axis] + 1, "TopK2: dynamic K is out of range");
            kVal = static_cast<int>(kTemp);
        }