I have a following problem:
I have two 1D arrays of the same length, for example:
data = array([5, 1, 1, 4, 2, 2, 1])
indexes = array([1, 4, 1, 2, 2, 4, 5])
I want to add elements from data that have the same value in the corresponding places in the indexes array, so in this example I want to get as the output:
expected_output = array([6, 3, 6, 6, 6, 3, 1])
I can compute this by:
lst = []
for i in indexes:
lst.append(data[indexes == i].sum())
output = np.array(lst)
The problem is this implementation is slow. When the input arrays have lengths on the order of one million this can take 10 minutes. How can I do this faster?
Additional information on my main problem
In general my problem is a sub-problem of an actual problem I'm trying to solve, namely how to divide all columns in a scipy sparse matrix by the sum of the columns (and then also divide all rows by the sum of rows), something I would have done in Pandas as:
df.div(df.sum(axis=0), axis=1)
and
df.div(df.sum(axis=1), axis=0)
But Pandas fails to handle the data sizes I'm working with. When I try to do the same thing using scipy csr_matrix I get division by 0 problems, because some rows and some columns have all zeros, but I guess I could get around that by setting certain values to 1 in the sum vectors. But the bigger problem is that code:
csr_arr / csr_arr.sum(axis=1)
results in
MemoryError: Unable to allocate 1.82 PiB for an array with shape (255999995, 999999) and data type int64
so the implementation tries to convert to a dense representation and this fails. I decided to implement this division by the sum vector entirely in the sparse representation, like so:
row_arr, col_arr = csr_arr.nonzero()
sum_col_lst = []
for i in row_arr:
sum_col_lst.append(csr_arr.data[row_arr == i].sum())
sum_row_lst = []
for i in col_arr:
sum_row_lst.append(csr_arr.data[col_arr == i].sum())
prob_a_given_b_data = csr_arr.data / np.array(sum_row_lst)
prob_b_given_a_data = csr_arr.data / np.array(sum_col_lst)
csr_prob_a_given_b = csr_matrix((prob_a_given_b_data, (row_arr, col_arr)), shape=shape)
csr_prob_b_given_a = csr_matrix((prob_b_given_a_data, (row_arr, col_arr)), shape=shape)
and this works fine, even for very large matrices (and I need to work with sparse matrices of shapes hundreds of millions times millions and holding millions of non-zero values) but like I said those for loops take a long time since nothing is vectorized there.