Суммируйте три элемента в каждой строке матрицы - без циклов - PullRequest
Новая
       9

Суммируйте три элемента в каждой строке матрицы - без циклов

0 голосов
/ 16 июня 2019

Я хочу написать функцию SumThree, которая получает матрицу размером (n x (n+2)) в виде массива и возвращает вектор v, который определен следующим образом: v[i] = A[i,i] + A[i,i+1] + A[i,i+2].

Ограничения: Нет петель, вы не можете использовать функцию sum (np.sum)

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

Примерно так (вот что я пробовал): 

def sumThree(array, curr=0, i=0, v=[]):
 if curr < array.size
   curr += array.shape[1]
 else
   return v
 count = array.size
 v[i] = array[curr,curr] + array[curr,curr+1] + array[curr,curr+2]
 sumThree(array, curr, i+1, v)

Я еще не пробовал, но думаю, что это должно работать: 

def sumThree(a):
  return (a.diagonal() + a.diagonal(1) + a.diagonal(2))

1 Ответ

1 голос
/ 16 июня 2019

Вы можете получить определенные диагонали (с требуемым смещением), используя numpy.diagonal

Не могли бы вы попробовать это и посмотреть, работает ли он для вас? 

def sumThree(a):
    return a.diagonal(offset=0) + a.diagonal(offset=1) + a.diagonal(offset=2)

вот полный пример: 

import numpy as np

def sumThree(a):
    return a.diagonal(offset=0) + a.diagonal(offset=1) + a.diagonal(offset=2)


my_arr = np.arange(48).reshape(6,8)
print(my_arr)
print("\n\n\n")
print(sumThree(my_arr))

выводит: 

[[ 0  1  2  3  4  5  6  7]
 [ 8  9 10 11 12 13 14 15]
 [16 17 18 19 20 21 22 23]
 [24 25 26 27 28 29 30 31]
 [32 33 34 35 36 37 38 39]
 [40 41 42 43 44 45 46 47]]




[  3  30  57  84 111 138]
...