Search code examples
pythonperformancesubstitutionsubmatrix

Substitute elements of a matrix at specific coordinates in python


I am trying to solve a "very simple" problem. Not so simple in Python. Given a large matrix A and another smaller matrix B I want to substitute certain elements of A with B. In Matlab is would look like this:

Given A, row_coord = [1,5,6] col_coord = [2,4], and a matrix B of size(3X2), A[row_coord, col_coord] = B

In Python I tried to use product(row_coord, col_coord) from the itertools to generate the set of all indexes that need to be accessible in A but it does not work. All examples on submatrix substitution refer to block-wise row_coord = col_coord examples. Nothing concrete except for the http://comments.gmane.org/gmane.comp.python.numeric.general/11912 seems to relate to the problem that I am facing and the code in the link does not work.

Note: I know that I can implement what I need via the double for-loop, but on my data such a loop adds 9 secs to the run of one iteration and I am looking for a faster way to implement this.

Any help will be greatly appreciated.


Solution

  • Assuming you're using numpy arrays then (in the case where your B is a scalar) the following code should work to assign the chosen elements to the value of B.

    itertools.product will create all of the coordinate pairs which we then convert into a numpy array and use in indexing your original array:

    import numpy as np
    from itertools import product
    
    A = np.zeros([20,20])
    
    col_coord = [0,1,3] 
    row_coord = [1,2] 
    
    coords = np.array(list(product(row_coord, col_coord)))
    
    B = 1
    
    A[coords[:,0], coords[:,1]] = B
    

    I used this excellent answer by unutbu to work out how to do the indexing.