在列表推导中创建 tmp 变量

Question

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

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 是一个在线函数，但是因为我在输出中多次需要它，所以我想暂时将它保存为一个变量，这样我就不必多次计算它了。但是，我不知道如何在列表理解期间执行此操作。

Answer 1

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

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)]]

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

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

Answer 2

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

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

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 列表理解快一点，这与人们告诉我的一致。

在列表推导中创建 tmp 变量

Creating tmp variables inside a list comprehension

python

list-comprehension

python-3.x