Как эффективно получить доступ / изменить один элемент в каждой строке матрицы

Question

У меня есть матрица A с размером (nr, nc), вектор индексов столбцов B (поэтому B имеет размер (nr, 1), и каждый элемент в B является целым числом от 1 до nc), и я хочу сделайте что-нибудь для каждого элемента в A, который имеет форму A (i, B (i)) для i между 1 и nr, ключевым фактором является эффективность. Для конкретности, скажем, C является вектором размера (nr, 1), цель состоит в том, чтобы сделать

for i=1:nr 
A(i,B(i))=A(i,B(i))+C(i)
end

более эффективно. Контекст обычно такой: nr >> nc (потому что когда nr большое, векторизация эффективна для многих операций). Я получил ускорение в 3 раза благодаря использованию функции индикатора:

for k=1:nc
A(:,k)=A(:,k)+(k==B).*C
end

Есть ли другие способы (более эффективные, надеюсь) сделать это? Я полагаю, что это похоже на многие вопросы о двойной индексации, но конкретно я сталкиваюсь с этим все время.

Answer 1

Использовать линейное индексирование:

idx = sub2ind(size(A), 1:nr, B');
A(idx) = A(idx) + C';

или (отредактированная версия с одним меньшим транспонированием)

idx = sub2ind(size(A), (1:nr)', B);
A(idx) = A(idx) + C;

Answer 2

Одним из способов будет использование линейного индексирования матрицы. Вам понадобится вектор v, содержащий смещения первого элемента в каждой строке, а затем индексирование с использованием A(v + B). Например:

>A=[1 2 3; 4 5 6; 7 8 9]

A =

 1     2     3
 4     5     6
 7     8     9

>B = [1 2 3] % we want the 1st element of row 1, 2nd of row 2, 3rd of row 3

>ii = [0 3 6] + B

>a(ii)

1     5     9

Примечание : Как показало groovingandi , также возможно (и более читабельно) использовать sub2ind для генерации вектора линейных индексов ii. Идея по сути та же самая.