python - 输出某个特定列的具有特定初始字符串的行

Question

我有一个制表符分隔的 txt 文件，如下所示:

A   B   aaaKP
C   D   bbbZ
E   F   cccLL

这是制表符分隔的。

如果

phrase = aaa or bbb
column = 3

然后我只想要第三列以 aaa 或 bbb 开头的行

输出将是

A   B   aaaKP
C   D   bbbZ

我有一个针对只有一个短语的情况的代码。

phrase, column = 'aaa', 3
fn = lambda l : len(l) >= column and len(l[column-1]) >= len(phrase) and phrase == l[column-1][:len(phrase)]
fp = open('output.txt', 'w')
fp.write(''.join(row for row in open('input.txt') if fn(row.split('\t'))))
fp.close()

但是如果有多个短语..我尝试过

phrase, column = {'aaa','bbb'}, 3

但是没有成功。

Answer 1

一般情况下，您可以使用带有分支的正则表达式来快速匹配和搜索:

import re

phrases = [ 'aaa', 'bbb' ]
column = 3

pattern = re.compile('|'.join(re.escape(i) for i in phrases))
column -= 1

with open('input.txt') as inf, open('output.txt', 'w') as outf:
    for line in inf:
        row = line.split('\t')
        if pattern.match(row[column]):
            outf.write(line)

代码根据所有可能的短语构建正则表达式，使用 re.escape 转义特殊字符。本例中的结果表达式为 aaa|bbb。 pattern.match 将字符串的开头与模式进行匹配(匹配必须从第一个字符开始)。

但是，如果您必须仅将字符串的开头与固定短语进行匹配，那么请注意 startswith接受一个元组，这是最快的代码:

 phrases = [ 'aaa', 'bbb' ]
 column = 3

 phrase_tuple = tuple(phrases)
 column -= 1

 with open('input.txt') as inf, open('output.txt', 'w') as outf:
     for line in inf:
         row = line.split('\t')
         if row[column].startswith(phrase_tuple):
             outf.write(line)

它还演示了如何使用上下文管理器打开文件，在output.txt之前打开input.txt，这样如果前者不存在，则后者不会被创建。最后表明，在没有任何生成器和 lambda 的情况下，这看起来最好。