在numpy数组中查找n维点

Question

我正在调查将点存储在 numpy 数组中是否有助于我搜索点，对此我有几个问题。

我有一个表示 3 维点的 Point 类。

class Point( object ):
  def __init__( self, x, y, z ):
    self.x = x
    self.y = y
    self.z = z

  def __repr__( self ):
    return "<Point (%r, %r, %r)>" % ( self.x, self.y, self.z )

我构建了一个 Point 对象列表。注意坐标(1, 2, 3)故意出现两次；这就是我要搜索的内容。

>>> points = [Point(1, 2, 3), Point(4, 5, 6), Point(1, 2, 3), Point(7, 8, 9)]

我将 Point 对象存储在一个 numpy 数组中。

>>> import numpy
>>> npoints = numpy.array( points )
>>> npoints
array([<Point (1, 2, 3)>, <Point (4, 5, 6)>, <Point (1, 2, 3)>,
   <Point (7, 8, 9)>], dtype=object)

我按照以下方式搜索坐标为(1, 2, 3)的所有点。

>>> numpy.where( npoints == Point(1, 2, 3) )
>>> (array([], dtype=int64),)

但是，结果没有用。因此，这似乎不是正确的方法。 numpy.where 是要使用的东西吗？是否有另一种方式来表达numpy.where 的条件会成功？

接下来我尝试将点的坐标存储在一个 numpy 数组中。

>>> npoints = numpy.array( [(p.x, p.y, p.z) for p in points ])
>>> npoints
array([[1, 2, 3],
      [4, 5, 6],
      [1, 2, 3],
      [7, 8, 9]])

我按照以下方式搜索坐标为(1,2,3)的所有点。

>>> numpy.where( npoints == [1,2,3] )
(array([0, 0, 0, 2, 2, 2]), array([0, 1, 2, 0, 1, 2]))

结果至少是我可以处理的。第一个返回值array([0, 0, 0, 2, 2, 2]) 中的行索引数组确实告诉我，我正在搜索的坐标位于npoints 的第0 行和第2 行。我可以做以下类似的事情。

>>> rows, cols = numpy.where( npoints == [1,2,3] )
>>> rows
array([0, 0, 0, 2, 2, 2])
>>> cols
array([0, 1, 2, 0, 1, 2])
>>> foundRows = set( rows )
>>> foundRows
set([0, 2])
>>> for r in foundRows:
...   # Do something with npoints[r]

但是，我觉得我并没有真正恰当地使用numpy.where，我只是在这种特殊情况下很幸运。

在 numpy 数组中查找所有出现的 n 维点（即具有特定值的行）的适当方法是什么？

保持数组的顺序很重要。

Answer 1

您可以在 Point 类中创建“丰富的比较”方法 object.__eq__(self, other) 以便能够在 Point 对象中使用 ==：

class Point( object ):
  def __init__( self, x, y, z ):
    self.x = x
    self.y = y
    self.z = z

  def __repr__( self ):
    return "<Point (%r, %r, %r)>" % ( self.x, self.y, self.z )
  def __eq__(self, other):
    return self.x == other.x and self.y == other.y and self.z == other.z

import numpy
points = [Point(1, 2, 3), Point(4, 5, 6), Point(1, 2, 3), Point(7, 8, 9)]
npoints = numpy.array( points )
found = numpy.where(npoints == Point(1, 2, 3))
print(found) # => (array([0, 2]),)