Эффективная замена элементов в массиве numpy

Предварительный ответ, который не работает

Вы можете легко векторизовать операцию обмена, используя массивы индексов (c1, r1, c2, r2) вместо перебора списков скалярных индексов.

c1 = np.array(<all the "c1" values>, dtype=int)
r1 = np.array(<all the "r1" values>, dtype=int)
c2 = np.array(<all the "c2" values>, dtype=int)
r2 = np.array(<all the "r2" values>, dtype=int)
A[c1, r1], A[c2, r2] = A[c2, r2], A[c1, r1]

Обратите внимание, что это выполняет все замены за один раз, что может отличаться от итеративного, если порядок замены имеет значение. По этой причине это неверный ответ на ваш вопрос.

Например. замена p1 на p2, затем p2 на p3 отличается от замены p1 и p2, а также p2 и p3 за один раз. В последнем случае и p1, и p3 получают исходное значение p2, а p2 получает последнее значение между p1 и p3, т. е. p3 (в соответствии с порядком их появления в массиве индексов). ).

Однако, поскольку вам нужна скорость, векторизация операции (каким-то образом) должна быть способом.

Добавление правильности к приведенному выше решению

Итак, как мы можем выполнить векторизованную подкачку, получая при этом нужный нам результат? Мы можем использовать гибридный подход, разбивая списки индексов на части (как можно меньше), где каждая часть содержит только уникальные точки, таким образом гарантируя, что порядок не имеет значения. Замена каждого фрагмента выполняется vercrorized, и мы перебираем только фрагменты.

Вот набросок того, как это может работать:

c1, r1 = np.array([ np.arange(10), np.arange(10) ])
c2, r2 = np.array([ [2,4,6,8,0,1,3,5,7,9], [9,1,6,8,2,2,2,2,7,0] ])
A = np.empty((15,15))

def get_chunk_boundry(c1, r1, c2, r2):
    a1 = c1 + 1j * r1
    a2 = c2 + 1j * r2
    set1 = set()
    set2 = set()
    for i, (x1, x2) in enumerate(zip(a1, a2)):
        if x1 in set2 or x2 in set1:
            return i
        set1.add(x1); set2.add(x2)
    return len(c1)

while len(c1) > 0:
    i = get_chunk_boundry(c1, r1, c2, r2)
    c1b = c1[:i]; r1b = r1[:i]; c2b = c2[:i]; r2b = r2[:i]
    print 'swapping %d elements' % i
    A[c1b,r1b], A[c2b,r2b] = A[c2b,r2b], A[c1b,r1b]
    c1 = c1[i:]; r1 = r1[i:]; c2 = c2[i:]; r2 = r2[i:]

Нарезка здесь будет быстрее, если вы сначала сохраните индексы в виде двумерного массива (N x 4).

Вот итеративное решение, которое я написал для справки (как один из способов справиться с возможными повторяющимися элементами):

def iter2(A, rc1, rc2):
    for r,c in zip(rc1.T, rc2.T):
        j,k = tuple(r),tuple(c)
        A[j],A[k] = A[k],A[j]
    return A

Например, если:

N = 4
Aref=np.arange(N)+np.arange(N)[:,None]*10

rc1=np.array([[0,0,0],[0,3,0]])
rc2=np.array([[3,3,2],[3,0,2]])

Тогда

 print(iter2(A.copy(), rc1,rc2))

Производит обмен углами, за которым следует обмен с (2,2):

[[22  1  2 30]
 [10 11 12 13]
 [20 21 33 23]
 [ 3 31 32  0]]

Решение shx2's, похоже, делает то же самое, хотя для моего тестового примера, похоже, есть ошибка в функции фрагментации.

Для массива shx2's test (15,15) мое итеративное решение работает в 4 раза быстрее! Он делает больше свопов, но меньше работы на каждый своп. Для больших массивов я ожидаю, что разбиение на фрагменты будет быстрее, но я не знаю, насколько больше нам придется идти. Это также будет зависеть от шаблона повторения.

Тупой векторизованный обмен с моими массивами:

A[tuple(rc1)],A[tuple(rc2)] = A[tuple(rc2)],A[tuple(rc1)]

В этом (15,15) примере это всего на 20% быстрее, чем мое итеративное решение. Очевидно, что нам нужен очень большой тестовый пример, чтобы получить серьезные тайминги.

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

def adapt_1d(A, rc1, rc2, fn=None):
    # adapt a multidim case to use a 1d iterator
    rc2f = np.ravel_multi_index(rc2, A.shape)
    rc1f = np.ravel_multi_index(rc1, A.shape)
    Af = A.flatten()
    if fn is None:
        for r,c in zip(rc1f,rc2f):
            Af[r],Af[c] = Af[c],Af[r]
    else:
        fn(Af, rc1f, rc2f)
    return Af.reshape(A.shape)