balanced_partition.m

function dim_ordered = balanced_partition(vals, M)
% Devide vals into M "balanced" subspaces
% This function is developed by Ge et al.
% See: "T. Ge, K. He, Q. Ke, and J. Sun. Optimized product quantization. IEEE Trans. on Pattern Analysis and Machine Intelligence, 36(4):744–755, 2014." 

dim = numel(vals);
dim_subspace = dim / M;

dim_tables = cell(M, 1);

fvals = log(vals+1e-20); %% balance the product of eigenvalues of the subspaces, i.e., the sum of log(eigenvalues)
fvals = fvals - min(fvals) + 1; %% make all positive

sum_list = zeros(M, 1);

current_subspaceIdx = 1;

for d=1:dim
    dim_tables{current_subspaceIdx} = [dim_tables{current_subspaceIdx}; d];
    
    sum_list(current_subspaceIdx) = sum_list(current_subspaceIdx) + fvals(d);
    
    if numel(dim_tables{current_subspaceIdx}) == dim_subspace %% this subspace is full
        sum_list(current_subspaceIdx) = 1e10;  %% do not use
    end
    
    [not_used, current_subspaceIdx] = min(sum_list);
end

dim_ordered = [];
for m=1:M
    dim_ordered = [dim_ordered; dim_tables{m}];
end

end