check if numpy array is subset of another array

2020-02-13 03:15发布

Similar questions have already been asked on SO, but they have more specific constraints and their answers don't apply to my question.

Generally speaking, what is the most pythonic way to determine if an arbitrary numpy array is a subset of another array? More specifically, I have a roughly 20000x3 array and I need to know the indices of the 1x3 elements that are entirely contained within a set. More generally, is there a more pythonic way of writing the following:

master=[12,155,179,234,670,981,1054,1209,1526,1667,1853] #some indices of interest
triangles=np.random.randint(2000,size=(20000,3)) #some data
for i,x in enumerate(triangles):
  if x[0] in master and x[1] in master and x[2] in master:
    print i

For my use case, I can safely assume that len(master) << 20000. (Consequently, it is also safe to assume that master is sorted because this is cheap).

标签: python numpy set
3条回答
手持菜刀,她持情操
2楼-- · 2020-02-13 04:10

You can do this easily via iterating over an array in list comprehension. A toy example is as follows:

import numpy as np
x = np.arange(30).reshape(10,3)
searchKey = [4,5,8]
x[[0,3,7],:] = searchKey
x

gives

 array([[ 4,  5,  8],
        [ 3,  4,  5],
        [ 6,  7,  8],
        [ 4,  5,  8],
        [12, 13, 14],
        [15, 16, 17],
        [18, 19, 20],
        [ 4,  5,  8],
        [24, 25, 26],
        [27, 28, 29]])

Now iterate over the elements:

ismember = [row==searchKey for row in x.tolist()]

The result is

[True, False, False, True, False, False, False, True, False, False]

You can modify it for being a subset as in your question:

searchKey = [2,4,10,5,8,9]  # Add more elements for testing
setSearchKey = set(searchKey)
ismember = [setSearchKey.issuperset(row) for row in x.tolist()]

If you need the indices, then use

np.where(ismember)[0]

It gives

array([0, 3, 7])
查看更多
Explosion°爆炸
3楼-- · 2020-02-13 04:10

A more natural (and possibly faster) solution for set operations in numpy is to use the functions in numpy.lib.arraysetops. These generally allow you to avoid having to convert back and forth between Python's set type. To check if one array is a subset of another, use numpy.setdiff1d() and test if the returned array has 0 length:

import numpy as np
a = np.arange(10)
b = np.array([1, 5, 9])
c = np.array([-5, 5, 9])
# is `a` a subset of `b`?
len(np.setdiff1d(a, b)) == 0 # gives False
# is `b` a subset of `a`?
len(np.setdiff1d(b, a)) == 0 # gives True
# is `c` a subset of `a`?
len(np.setdiff1d(c, a)) == 0 # gives False

You can also optionally set assume_unique=True for a potential speed boost.

I'm actually a bit surprised that numpy doesn't have something like a built-in issubset() function to do the above (analogous to set.issubset()).

Another option is to use numpy.in1d() (see https://stackoverflow.com/a/37262010/2020363)

Edit: I just realized that at some point in the distant past this bothered me enough that I wrote my own simple function:

def issubset(a, b):
    """Return whether sequence `a` is a subset of sequence `b`"""
    return len(np.setdiff1d(a, b)) == 0
查看更多
家丑人穷心不美
4楼-- · 2020-02-13 04:12

Here are two approaches you could try:

1, Use sets. Sets are implemented much like python dictionaries and have have constant time lookups. That would look much like the code you already have, just create a set from master:

master = [12,155,179,234,670,981,1054,1209,1526,1667,1853]
master_set = set(master)
triangles = np.random.randint(2000,size=(20000,3)) #some data
for i, x in enumerate(triangles):
  if master_set.issuperset(x):
    print i

2, Use search sorted. This is nice because it doesn't require you to use hashable types and uses numpy builtins. searchsorted is log(N) in the size of master and O(N) in the size of triangels so it should also be pretty fast, maybe faster depending on the size of your arrays and such.

master = [12,155,179,234,670,981,1054,1209,1526,1667,1853]
master = np.asarray(master)
triangles = np.random.randint(2000,size=(20000,3)) #some data
idx = master.searchsorted(triangles)
idx.clip(max=len(master) - 1, out=idx)
print np.where(np.all(triangles == master[idx], axis=1))

This second case assumes master is sorted, as searchsorted implies.

查看更多
登录 后发表回答