function [codistrL, codistrI] = hGetCodistForTopK(codistrA, k, sortDim, distDim)
%HGETCODISTFORTOPK creates new codistributors for the output of distributed
%maxk, mink, and topkrows
%   Input   k         Number of elements to return in sort dimension
%           sortDim   Dimension on which sorting is done
%           distDim   Dimension on which the array is distributed
%   Output  codistrL  Codistributor for output top k elements
%           codistrI  Codistributor for output indices from topkrows

%   Copyright 2019-2022 The MathWorks, Inc.

globalSizeL = codistrA.Cached.GlobalSize;
globalSizeL(sortDim) = k;
if sortDim == distDim
    partSum = min(cumsum(codistrA.Partition), k);
    PartitionL = diff([0 partSum]);
    PartitionI = PartitionL;
else
    PartitionL = codistrA.Partition;
    PartitionI = [1 zeros(1, spmdSize-1)];
end	
codistrL = codistributor1d(distDim, PartitionL, globalSizeL);
if nargout == 2
    globalSizeL(2) = 1;
    codistrI = codistributor1d(distDim, PartitionI, globalSizeL);
end
