在numpy中以广度优先顺序对树的节点进行排序

Question

假设我有以下numpy 数组：

arr = numpy.array([[1, 7], [2, 0], [2, 1], [2, 3], [3, 4], [3, 5], [5, 6]])

假设我选择了一个特定的键，在本例中为 2。然后我希望排序如下：

arr_sorted = [[2, 0], [2, 1], [2, 3], [1, 7], [3, 4], [3, 5], [5, 6]]

这个想法是首先从所有key为2的元素开始，然后去那些key是前一个key的值的条目。

从 2 开始，条目为 [2, 0], [2, 1], [2, 3]。因此，2 之后的下一个键将是 0、1、3。没有以 0 作为键的条目。以 1 为键的条目有一个：[1, 7]。有两个条目以 3 为键：[3, 4], [3, 5]。下一个未处理的键是7，但它没有条目。 4 也是如此。有一个条目以 5 为键：[5, 6]。 6 没有条目。

是否有任何numpy 或dictionary 技巧来实现这一点？

我最近的尝试如下：

def bfs_finder(d, start):
  queue = deque([start])
  seen = [start]
  results = []
  while queue:
    _vertices = queue.popleft()
    current = [i for i, a in enumerate(d) if len([x for x in a if x in _vertices])==1 and i not in seen]
    curr1 = [a[1] for i, a in enumerate(d) if len([x for x in a if x in _vertices]) == 1 and i not in seen]
    if len(current)>0:
        results.extend(curr1)
        queue.extend(curr1)
        seen.extend(current)
  return results

但是，我实际上收到了current = [i for i, a in enumerate(d) if len([x for x in a if x in _vertices])==1 and i not in seen] TypeError: argument of type 'int' is not iterable 的错误。任何有关如何修复此错误的建议，以及是否有任何好的改进，我们将不胜感激。

Answer 1

您可以使用set 来保存已处理的键，并使用deque 或类似的堆栈兼容容器来保存已处理的值（列表也可以）。由于您有 numpy 数组，因此您可以按行对数组进行预排序，并在第一列使用 np.searchsorted 从中提取条带。

算法是这样的：

1. 按行对数组进行排序
2. 预分配输出数组
3. 将密钥添加到stack
4. 当堆栈不为空时
   5. 从堆栈中弹出一个键
   6. 如果密钥在集合中，丢弃并继续
   7. 将密钥添加到集合中
   8. 使用二分法在第一列（开始和结束索引）中查找键
   9. 如果存在键，则将 swath 复制到输出数组

排序操作是O(n log(n))。二分算法也是如此（大约n 搜索每个O(log(n))）。因此，您的总算法复杂度不应超过 O(n log(n))，这对于奇怪的排序算法来说非常好。

根据this answer，使用np.argsort 按行排序似乎是最快的，幸运的是np.searchsorted 接受sorter 参数。

这是一个示例实现：

import numpy as np
from collections import deque

def bfs_finder(d, start):
    sorter = np.argsort(d[:, 0])
    done = set()
    todo = deque([start])
    output = np.empty_like(d)
    pos = 0
    while todo:
        key = todo.popleft()
        if key in done:
            continue
        done.add(key)
        left = np.searchsorted(d[:, 0], key, 'left', sorter)
        if left >= d.shape[0] or d[sorter[left], 0] != key:
            continue
        right = np.searchsorted(d[:, 0], key, 'right', sorter)
        next = pos + right - left
        output[pos:next, :] = d[sorter[left:right], :]
        todo.extend(output[pos:next, 1])
        pos = next
    return output

arr = np.array([[1, 7], [2, 0], [2, 1], [2, 3], [3, 4], [3, 5], [5, 6]])
print(bfs_finder(arr, 2))

IDEOne Link

[[2 0]
 [2 1]
 [2 3]
 [1 7]
 [3 4]
 [3 5]
 [5 6]]

此解决方案假定您不会在输出中丢失任何原始键。如果遇到该问题，请从已处理键的集合中减去第一列的集合，然后决定如何处理余数。

您可以通过将整个堆栈传递给searchsorted 的每次调用而不是一次执行一个元素来获得一些额外的里程。