根据下标获取行平均值的快速方法

Question

我有一个数据，可以通过以下方式模拟：

N = 10^6;%10^8;
K = 10^4;%10^6; 

subs = randi([1 K],N,1);
M = [randn(N,5) subs];
M(M<-1.2) = nan;

换句话说，它是一个矩阵，其中最后一行是下标。 现在我想为每个下标计算nanmean()。我还想为每个下标保存行数。我有一个“虚拟”代码：

uniqueSubs = unique(M(:,6));
avM = nan(numel(uniqueSubs),6);
for iSub = 1:numel(uniqueSubs)
    tmpM = M(M(:,6)==uniqueSubs(iSub),1:5);
    avM(iSub,:) = [nanmean(tmpM,1) size(tmpM,1)];
end

问题是，它太慢了。我希望它适用于 N = 10^8 和 K = 10^6（请参阅这些变量定义中的注释部分。

如何更快地找到数据的平均值？

Answer 1

这对于findgroups 和splitapply 来说听起来是个完美的工作。

% Find groups in the final column
G = findgroups(M(:,6));
% function to apply per group
fcn = @(group) [mean(group, 1, 'omitnan'), size(group, 1)];
% Use splitapply to apply fcn to each group in M(:,1:5)
result = splitapply(fcn, M(:, 1:5), G);
% Check
assert(isequaln(result, avM));

Answer 2

M = sortrows(M,6); % sort the data per subscript
IDX = diff(M(:,6)); % find where the subscript changes
tmp = find(IDX);
tmp = [0 ;tmp;size(M,1)]; % add start and end of data
for iSub= 2:numel(tmp)
    % Calculate the mean over just a single subscript, store in iSub-1
    avM2(iSub-1,:) = [nanmean(M(tmp(iSub-1)+1:tmp(iSub),1:5),1) tmp(iSub)-tmp(iSub-1)];tmp(iSub-1)];
end

这比您在我的计算机上的原始代码快大约 60 倍。加速主要来自对数据进行预排序，然后找到下标发生变化的所有位置。这样，您不必每次都遍历整个数组来找到正确的下标，而只需检查每次迭代所需的内容。因此，您可以计算大约 100 行的平均值，而不必首先检查 1,000,000 行是否需要该迭代。

因此：在原来的你检查numel(uniqueSubs)，在这种情况下是10,000，是否所有N，这里是1,000,000，数字属于某个类别，结果是10^12次检查。建议的代码对行进行排序（排序为NlogN，因此此处为 6,000,000），然后在整个数组上循环一次，无需额外检查。

---

为了补全，这里是原始代码，以及我的版本，它表明两者是相同的：

N = 10^6;%10^8;
K = 10^4;%10^6; 

subs = randi([1 K],N,1);
M = [randn(N,5) subs];
M(M<-1.2) = nan;

uniqueSubs = unique(M(:,6));
%% zlon's original code
avM = nan(numel(uniqueSubs),7); % add the subscript for comparison later
tic
uniqueSubs = unique(M(:,6));
for iSub = 1:numel(uniqueSubs)
    tmpM = M(M(:,6)==uniqueSubs(iSub),1:5);
    avM(iSub,:) = [nanmean(tmpM,1) size(tmpM,1) uniqueSubs(iSub)];
end
toc
%%%%% End of zlon's code
avM = sortrows(avM,7); % Sort for comparison

%% Start of Adriaan's code
avM2 = nan(numel(uniqueSubs),6);
tic
M = sortrows(M,6);
IDX = diff(M(:,6));
tmp = find(IDX);
tmp = [0 ;tmp;size(M,1)];
for iSub = 2:numel(tmp)
    avM2(iSub-1,:) = [nanmean(M(tmp(iSub-1)+1:tmp(iSub),1:5),1) tmp(iSub)-tmp(iSub-1)];
end
toc %tic/toc should not be used for accurate timing, this is just for order of magnitude
%%%% End of Adriaan's code

all(avM(:,1:6) == avM2) % Do the comparison
% End of script

% Output
Elapsed time is 58.561347 seconds.
Elapsed time is 0.843124 seconds. % ~70 times faster

ans =

  1×6 logical array

   1   1   1   1   1   1 % i.e. the matrices are equal to one another