时间戳

首页 > 解决方案 > 在 python 中折叠嵌套列表以容纳可变数量的列表（可能使用 itertools）

python - 在 python 中折叠嵌套列表以容纳可变数量的列表（可能使用 itertools）

问题描述

我正在尝试扩展下面的函数，它需要 2 个输入列表并使用它们来填充输出列表，这样每个输入列表都是一个嵌套函数：

import numpy as np

## condition lists 
cond_list_a = [1,0,1,0,1,0,1,0]
cond_list_b = [2,2,2,2,4,4,4,4]

## elements to populate output list 
elem_pool = [3,5]

## function that takes 2 condition lists 
def list_builder (cond_list_a, cond_list_b, elem_pool):

    trials = len(cond_list_a)                              ## output length
    output_list = [0]*trials                               ## creating a base list 

    for elem_a in np.unique(cond_list_a):                  ## for each unique element in condition list a 
        for elem_b in np.unique(cond_list_b):              ## for each unique eleement in condition list b 
            pool_idx = 0                                   ## start the index to draw elements from at 0 
            for i in range(trials):                        
                if (cond_list_a[i] == elem_a) and (cond_list_b[i] == elem_b): ## for each combination 
                    output_list[i] = elem_pool[pool_idx]   ## populate output list
                    pool_idx += 1                          ## update pool index
    return output_list

## run function
list_builder(cond_list_a, cond_list_b, elem_pool)

这将产生以下 output_list = [3, 3, 5, 5, 3, 3, 5, 5]

我的目标是构建一个可以接收任意数量的输入列表的函数，以同样的方式填充输出列表。一段时间以来，我一直在尝试使用 itertools.product，它可以创建所有可能的输入元素组合，但我很难做到这一点。我会很感激任何想法。不需要使用 itertools。

这是我（失败的）尝试之一的示例：

import itertools as it

## list of condition lists: 
cond_lists = [cond_list_a, cond_list_b]

## function that could take handle any number of condition lists 
def better_list_builder(condition_lists, elem_pool):

    trials = len(condition_lists[0])                            ## output length
    ouput_list = [0]*trials                                     ## creating a base list 
    
    cond_list_combos = list(it.product(*condition_lists))       ## get all the list combos
    uniq_combos = list(set(cond_list_combos))                   ## pare them down to the unique element combinations
    
    for each_combo in uniq_combos:                              ## for each combination of elements
        for j in range(len(condition_lists)):                   ## for each input condition list
            pool_idx = 0                                        ## start the index to draw elements from at 0        
            for i in range(trials):  
                if condition_lists[j][i] == each_combo[j]:      ## if the condition combo matches that list entry
                    output_list[i] = elem_pool[pool_idx]        ## populate output list
                    pool_idx += 1                               ## update pool index
    
    return output_list

## run function                    
better_list_builder(condition_lists, elem_pool)

标签： pythonfunctionitertoolsnested-lists

解决方案

知道了！

def list_builder (cond_list_master, elem_pool):

    trials = len(cond_list_master[0])                  ## output length
    output_list = [0]*trials                           ## creating a base list 

    somelist = []
    for idx in range(len(cond_list_master)):
        somelist.append(np.unique(cond_list_master[idx]))

    for element in itertools.product(*somelist):  # element is what you called (elem_a, elem_b, elem_c etc. )
        pool_idx = 0                                   ## start the index to draw elements from at 0 
        for i in range(trials): 
            flag = True
            for idx in range(len(cond_list_master)):
                flag = flag and (element[idx] == cond_list_master[idx][i])
            if flag:
                output_list[i] = elem_pool[pool_idx]   ## populate output list
                pool_idx += 1                          ## update pool index
    return output_list

推荐阅读