function x = KroSupThresholdOperator(x, N, s)
X = reshape(x, N.');   % Size: N1 x N2 x ... NI
I = length(N);
for dims = I:-1:1
    X_unfold = tens2mat(X,dims);
    [~, idx] = maxk(sum(abs(X_unfold).^2,2), s(dims));
    X_temp = zeros(size(X_unfold));
    X_temp(idx,:) = X_unfold(idx,:);

    X = fold(X_temp, dims, N);
end

x = vec(X);
end

% fold: fold matrix back to tensor
function T = fold(X_mat, mode, sz)
order = [mode, 1:mode-1, mode+1:numel(sz)];
perm_order = [2:mode,1,mode+1:numel(sz)];
dim_temp = order;
dim_temp((order==mode)) = [];
T = permute(reshape(X_mat, [sz(mode), sz(dim_temp).']), perm_order);
end