在列表推导中创建 tmp 变量

Creating tmp variables inside a list comprehension

我有以下列表理解,我想将其优化并变成一个单一的列表理解,但我不知道如何处理这些转变:

lines = [line for line in lines if len(line) > 55]
shifts = [int(int(line.split()[1]) >= 1000000)+int(int(line.split()[4]) >= 100000) for line in lines]
xyz = [(float(line[30+s:38+s]),float(line[38+s:46+s]),float(line[46+s:54+s])) for s,line in zip(shifts,lines)]

我知道组合起来应该是这样的:

xyz = [(float(line[30 + s:38 + s]), float(line[38 + s:46 + s]), float(line[46 + s:54 + s])) for line in lines if len(line) > 55]

但我仍然需要以某种方式 add/define s 变量。我怀疑我可以为此使用海象运算符,但实际上我不确定,因为我真的不想测试 condition ,而只是想分配它。 所以下面的例子不起作用,因为 s 被用作条件,有时 s == 0,这意味着它正在删除移位为零的示例,这不是我想要的:

xyz = [(float(line[30 + s:38 + s]), float(line[38 + s:46 + s]), float(line[46 + s:54 + s])) for line in lines if len(line) > 55 and (s := int(int(line.split()[1]) >= 1000000)+int(int(line.split()[4]) >= 100000))]

我当然可以只在所有位置使用 s 的定义而不是 s,但这看起来很丑陋且效率低下。那么有没有更好的方法呢?

编辑: 我需要这段代码尽可能快,这就是为什么我使用列表理解形式而不是循环形式,也是为什么我想将 3 个列表理解合并为一个。

循环形式的类似代码如下所示:

for line in lines:
    shift = 0
    dat = line.strip('\n')
    data = dat.split()
    if len(dat) > 55:
        if int(data[4]) >= 100000:
            shift += 1
        if int(data[1]) >= 1000000:
            shift += 1
        x = float(dat[30+shift:38+shift])
        y = float(dat[38+shift:46+shift])
        z = float(dat[46+shift:54+shift])
        X.append(x)
        Y.append(y)
        Z.append(z)
    else:
        continue

我知道我的例子不是最容易理解的,但我想要的基本上是以下理解形式

for line in lines:
    s = g(line)
    result = (f1(line,s),f2(line,s),f3(line,s))

其中g,f1,f2,f3是一些不重要的函数。所以最重要的是 s 是一个在线函数,但是因为我在输出中多次需要它,所以我想暂时将它保存为一个变量,这样我就不必多次计算它了。但是,我不知道如何在列表理解期间执行此操作。

CPython 3.9甚至优化为简单赋值的"idiom for assignment a temporary variable in comprehensions",用于splits:

xyz = [(float(line[30+s:38+s]),
        float(line[38+s:46+s]),
        float(line[46+s:54+s]))
       for line in lines
       if len(line) > 55
       for split in [line.split()]
       for s in [(int(split[1]) >= 1000000) +
                 (int(split[4]) >= 100000)]]

在我看来,像这样分布在多条短线上的整个内容非常易读。

顺便说一句,我摆脱了显式的 boolint 转换。当您添加两个 bool 时,您无论如何都会得到一个 int(例如,True + True2)并且它是 much faster(尽管在您的整体代码中可能并不重要)。

更新昨天发的结果今天不能复现,不知道昨天哪里出了问题,但是这里是新的结果,也和什么一致凯莉得到了。

根据不同的建议我做了一个对比:

import numpy as np
import time
def get_origo_and_basis_vectors_old(file_pdb_in):
    with open(file_pdb_in, 'r') as f:
        lines = f.readlines()[1:]
    xyz = []
    XYZapp = xyz.append
    for line in lines:
        data = line.split()
        if len(line) > 55:
            shift = (int(data[4]) >= 100000) + (int(data[1]) >= 1000000)
            XYZapp((float(line[30+shift:38+shift]),float(line[38+shift:46+shift]),float(line[46+shift:54+shift])))
        else:
            continue
    XYZ = np.asarray(xyz)
    origo = np.mean(XYZ,axis=0)
    basisvectors = np.max(XYZ,axis=0) - np.min(XYZ,axis=0)
    return origo, basisvectors


def get_origo_and_basis_vectors(file_pdb_in):
    with open(file_pdb_in, 'r') as f:
        lines = f.readlines()[1:]
    lines = [line for line in lines if len(line) > 55]
    shifts = [(int(line.split()[1]) >= 1000000)+(int(line.split()[4]) >= 100000) for line in lines]
    xyz = [(float(line[30+s:38+s]),float(line[38+s:46+s]),float(line[46+s:54+s])) for s,line in zip(shifts,lines)]
    XYZ = np.asarray(xyz)
    origo = np.mean(XYZ,axis=0)
    basisvectors = np.max(XYZ,axis=0) - np.min(XYZ,axis=0)
    return origo, basisvectors



def get_origo_and_basis_vectors_new(file_pdb_in):
    with open(file_pdb_in, 'r') as f:
        lines = f.readlines()[1:]
    xyz = [(float(line[30 + s:38 + s]),
            float(line[38 + s:46 + s]),
            float(line[46 + s:54 + s]))
           for line in lines
           if len(line) > 55
           for split in [line.split()]
           for s in [(int(split[1]) >= 1000000) +
                     (int(split[4]) >= 100000)]]
    XYZ = np.asarray(xyz)
    origo = np.mean(XYZ,axis=0)
    basisvectors = np.max(XYZ,axis=0) - np.min(XYZ,axis=0)
    return origo, basisvectors

if __name__=='__main__':
    file = '/media/tue/Data/Data/test_mini/relax_pdb/calc/AF_1W_1WU_1WUZ_1_A/leap/AF_1W_1WU_1WUZ_1_A_neutral.pdb'

    t0=time.time()
    for i in range(10):
        oo2, aa2 = get_origo_and_basis_vectors_new(file)
    print(f"1 list comprehension {time.time()-t0:2.2f}s")


    t0=time.time()
    for i in range(10):
        oo, aa = get_origo_and_basis_vectors(file)
    print(f"3 list comprehensions {time.time()-t0:2.2f}s")

    t0=time.time()
    for i in range(10):
        oo3, aa3 = get_origo_and_basis_vectors_old(file)
    print(f"for loop {time.time()-t0:2.2f}s")

我从中得到以下时间:

1 list comprehension 11.81s
3 list comprehensions 12.71s
for loop 11.87s

所以 1 列表理解和 for 循环似乎都比 3 列表理解快一点,这与人们告诉我的一致。