相当于 R "split"函数的 Python

Question

在 R 中，您可以根据另一个向量的因子拆分一个向量:

> a <- 1:10
  [1]  1  2  3  4  5  6  7  8  9 10
> b <- rep(1:2,5)
  [1] 1 2 1 2 1 2 1 2 1 2

> split(a,b)

   $`1`
   [1] 1 3 5 7 9
   $`2`
   [1]  2  4  6  8 10

因此，根据另一个列表的值(根据因子的顺序)对一个列表(在 python 方面)进行分组。

除了 itertools.groupby 方法之外，在 python 中有什么方便的方法吗？

Answer 1

从您的示例来看，b 中的每个元素看起来都包含将存储节点的 1 索引列表。 Python 缺少 R 似乎拥有的自动数值变量，因此我们将返回一个列表元组。如果您可以执行零索引列表，并且您只需要两个列表(即，对于您的 R 用例，1 和 2 是唯一的值，在 Python 中它们将是 0 和 1)

>>> a = range(1, 11)
>>> b = [0,1] * 5

>>> split(a, b)
([1, 3, 5, 7, 9], [2, 4, 6, 8, 10])

然后你可以使用itertools.compress:

def split(x, f):
    return list(itertools.compress(x, f)), list(itertools.compress(x, (not i for i in f)))

如果您需要更一般的输入(多个数字)，类似下面的内容将返回一个 n 元组:

def split(x, f):
    count = max(f) + 1
    return tuple( list(itertools.compress(x, (el == i for el in f))) for i in xrange(count) )  

>>> split([1,2,3,4,5,6,7,8,9,10], [0,1,1,0,2,3,4,0,1,2])
([1, 4, 8], [2, 3, 9], [5, 10], [6], [7])