如何用Python高效将SELFIES字符串分词为列表？

我正在使用SELFIES（自引用嵌入字符串），它是分子的一种字符串表示形式，由方括号包裹的一系列token组成，比如丙烷可表示为 "[C][C][C]"。我的需求是将这类字符串高效转换为对应的token列表，示例如下：

selfies = "[C][C][C]"
tokens = some_function(selfies)
tokens  # 输出：["[C]","[C]","[C]"]

我测试了多种实现方法，以下是详细对比：

方法1：SELFIES原生分词函数 #

使用SELFIES库自带的分词函数，该函数会处理SELFIES中的分隔符.，但我的场景中不会出现该字符。

def split_selfies(selfies: str) -> Iterator[str]:
    """Tokenizes a SELFIES string into its individual symbols.
    :param selfies: a SELFIES string.
    :return: the symbols of the SELFIES string one-by-one with order preserved.
    :Example:
    >>> import selfies as sf
    >>> list(sf.split_selfies("[C][=C][F].[C]"))
    ['[C]', '[=C]', '[F]', '.', '[C]']
    """

    left_idx = selfies.find("[")

    while 0 <= left_idx < len(selfies):
        right_idx = selfies.find("]", left_idx + 1)
        if right_idx == -1:
            raise ValueError("malformed SELFIES string, hanging '[' bracket")

        next_symbol = selfies[left_idx: right_idx + 1]
        yield next_symbol

        left_idx = right_idx + 1
        if selfies[left_idx: left_idx + 1] == ".":
            yield "."
            left_idx += 1

性能测试：

%%timeit
tokens = list(sf.split_selfies(selfies))
# 3.41 µs ± 22.7 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)

该方法速度较慢，推测是因为需要将生成器转换为列表。

方法2：库开发者实现的分词函数 #

这是SELFIES相关项目中开发者提供的实现，针对无.的场景优化。

def get_selfie_chars(selfies):
    '''Obtain a list of all selfie characters in string selfie
    
    Parameters: 
    selfie (string) : A selfie string - representing a molecule 
    
    Example: 
    >>> get_selfie_chars('[C][=C][C][=C][C][=C][Ring1][Branch1_1]')
    ['[C]', '[=C]', '[C]', '[=C]', '[C]', '[=C]', '[Ring1]', '[Branch1_1]']
    
    Returns:
    chars_selfie: list of selfie characters present in molecule selfie
    '''
    chars_selfie = []  # A list of all SELFIE sybols from string selfie
    while selfies != '':
        chars_selfie.append(selfies[selfies.find('['): selfies.find(']')+1])
        selfies = selfies[selfies.find(']')+1:]
    return chars_selfie

性能测试：

%%timeit
tokens = get_selfie_chars(selfies)
# 3.44 µs ± 43.9 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)

耗时和原生函数基本持平。

方法3：自定义列表推导+split实现 #

我自己结合列表推导、切片和split()实现的方法，针对无.的场景优化。

def selfies_split(selfies):
    return [block+"]" for block in selfies.split("]")][:-1]

性能测试：

%%timeit
tokens = selfies_split(selfies)
# 1.05 µs ± 53.2 ns per loop (mean ± std. dev. of 7 runs, 1000000 loops each)

速度比前两种方法快约3倍。

补充方法：社区建议的三种实现 #

方法4：字符串替换+split #

def stackoverflow_1_split(selfies):
    atoms = selfies[1:-1].replace('][', "$").split("$")
    return list(map('[{}]'.format, atoms))

性能测试：

%%timeit
tokens = stackoverflow_1_split(selfies)
# 1.75 µs ± 101 ns per loop (mean ± std. dev. of 7 runs, 1000000 loops each)

若不转换为列表，该方法速度更快，但列表是必需输出格式。

方法5：正则表达式匹配 #

import re
def stackoverflow_2_split(selfies):
    return re.findall(r".*?]", selfies)

性能测试：

%%timeit
tokens = stackoverflow_2_split(selfies)
# 1.81 µs ± 110 ns per loop (mean ± std. dev. of 7 runs, 1000000 loops each)

正则表达式方案并未优于其他方法，性能不如自定义的方法3。

方法6：字符串替换+空格分割（目前最快） #

def stackoverflow_3_split(selfies):
    return selfies.replace(']', '] ').split()

性能测试：

%%timeit
tokens = stackoverflow_3_split(selfies)
# 485 ns ± 4.04 ns per loop (mean ± std. dev. of 7 runs, 1000000 loops each)

这是目前测试中最快的方案，速度约为方法3的2倍。

内容的提问来源于stack exchange，提问作者Etienne Reboul