Python等效于R“拆分”功能

提问

在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中还有什么方便的方法吗?

最佳答案

在您的示例中,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])