CODEWITHSUNDEEP
pythonarraysnumpy
kevinkayaks
kevinkayaks

Reputation: 2726

Finding mutual elements between two higher-dimensional numpy arrays

I am new to python. Here is my problem. 

import numpy as np

def neighbors(indexset, i,j):
    temp = np.array([[i-1,j],[i+1,j],[i,j-1],[i,j+1]])
    for ele in temp:
        if ele in indexset: 
            print(ele)

indexset = np.array([[0,1],[1,1],[2,1],[3,1]])
neighbors(indexset, 0,0)

When I run this I get values I don't understand

neighbors(indexset, 0,0)
[1 0]
[ 0 -1]
[0 1]

What am I doing wrong? Why doesn't this return only [0,1]?

Upvotes: 1

Views: 156

Answers (1)

Angus Williams
Angus Williams

Reputation: 2334

I think you are getting stange results because ele in temp is only appropriate if ele is a scalar. It uses the numpy function __contains__, which is equivalent to (a==b).any(). If you do this with python lists instead of numpy arrays, it works:

def neighbors(indexset, i,j):
    temp = [[i-1,j],[i+1,j],[i,j-1],[i,j+1]]
    for ele in temp:
        if ele in indexset: 
            print(ele)

indexset = [[0,1],[1,1],[2,1],[3,1]]
neighbors(indexset, 0,0)

Will print [0,1] as expected. If indexset is always a numpy array, you can use tolist:

import numpy as np

def neighbors(indexset, i,j):
    temp = [[i-1,j],[i+1,j],[i,j-1],[i,j+1]]
    for ele in temp:
        if ele in indexset.tolist(): 
            print(ele)

indexset = np.array([[0,1],[1,1],[2,1],[3,1]])
neighbors(indexset, 0,0)

Upvotes: 2

Related Questions