SMILES字符串索引：一阶括号逻辑修复及二阶括号适配需求

需求说明 #

基础索引规则 #

• 非括号分支内的原子标记为0
• 一级括号内的原子按顺序从1开始计数
• 嵌套括号（括号层级>1）内的原子继承最近的一级括号分支的计数

示例：

CCC(C)CC → [0 0 0 1 0 0]
CC(CCC)CC → [0 0 1 2 3 0 0]
CC(C)(C)C → [0 0 1 1 0]
CC(CCC)(C)C → [0 0 1 2 3 1 0]
C(C(C)C)C → [0 1 1 2 0]

现有代码问题 #

对于SMILES字符串CCC(CC(C)CC)(C)C，期望输出为[0, 0, 0, 1, 2, 2, 3, 4, 1, 0]，但现有代码输出为[0, 0, 0, 1, 2, 2, 3, 4, 4, 0]，错误原因是处理完嵌套分支后未正确重置计数器状态，导致后续独立分支继承了嵌套内的计数。

扩展需求 #

支持指定括号阶数（如二阶括号((...))），仅对该阶数括号内的原子进行计数，其余原子标记为0，计数逻辑与一阶规则一致：目标阶数括号内原子顺序计数，嵌套层级超过目标阶数的原子继承最近的目标阶数分支的计数。

现有代码 #

import pandas as pd
from rdkit import Chem

def smile_grouping(s):
    group_counter = 1
    res = []
    open_brackets = 0
    branch_start_index = None
    last_non_nested_group = None

    for i, letter in enumerate(s):
        if letter == '(':
            open_brackets += 1
            if open_brackets == 1:
                branch_start_index = i
                if last_non_nested_group is not None:
                    group_counter = last_non_nested_group + 1
        elif letter == ')':
            if open_brackets == 1:
                last_non_nested_group = None
            open_brackets -= 1
            if open_brackets == 0:
                branch_start_index = None
        elif letter not in ['[', ']', '+', '-', '0', '1', '2', '3', '4', '5', '6', '7', '8', '9', 'i', 'e']:
            if open_brackets == 1:
                if branch_start_index is not None and branch_start_index + 1 != i:
                    group_counter += 1
                res.append(group_counter)
                last_non_nested_group = group_counter
            elif open_brackets == 0:
                res.append(0)
            elif open_brackets > 1:
                res.append(last_non_nested_group)

    mol = Chem.MolFromSmiles(s)
    num_atoms = mol.GetNumAtoms()
    while len(res) < num_atoms:
        res.append(0)

    return res

df = pd.DataFrame(
    {'SMILES': [
        "CCC(CC(C)CC)(C)C",
        "CCC[I+](C)(C)C",
        "CCC(CCC(C))C"
    ]})
df['Indexed SMILES'] = df['SMILES'].apply(smile_grouping)
print(df)

修正与扩展后的代码 #

import pandas as pd
from rdkit import Chem

def smile_grouping(s, target_depth=1):
    group_counter = 1
    res = []
    open_brackets = 0
    # 记录各目标层级分支的最后计数，处理多分支场景
    branch_last_counts = []
    current_target_count = None

    for i, letter in enumerate(s):
        if letter == '(':
            open_brackets += 1
            # 进入目标层级括号，初始化计数器
            if open_brackets == target_depth:
                # 如果有之前的目标分支计数，新分支从它+1开始
                group_counter = branch_last_counts[-1] + 1 if branch_last_counts else 1
                current_target_count = group_counter
        elif letter == ')':
            # 退出目标层级括号，记录该分支的最后计数
            if open_brackets == target_depth:
                if current_target_count is not None:
                    branch_last_counts.append(current_target_count)
                current_target_count = None
            open_brackets -= 1
            # 回到非目标层级，清空当前目标计数
            if open_brackets < target_depth:
                current_target_count = None
        elif letter not in ['[', ']', '+', '-', '0', '1', '2', '3', '4', '5', '6', '7', '8', '9', 'i', 'e']:
            # 目标层级内的原子：计数并更新当前目标计数
            if open_brackets == target_depth:
                res.append(group_counter)
                current_target_count = group_counter
                group_counter += 1
            # 嵌套层级超过目标：继承最近的目标分支计数
            elif open_brackets > target_depth:
                res.append(current_target_count if current_target_count is not None else 0)
            # 非目标层级：标记为0
            else:
                res.append(0)

    # 对齐RDKit解析的原子数量
    mol = Chem.MolFromSmiles(s)
    num_atoms = mol.GetNumAtoms()
    while len(res) < num_atoms:
        res.append(0)

    return res

# 测试示例
df = pd.DataFrame(
    {'SMILES': [
        "CCC(CC(C)CC)(C)C",
        "CCC[I+](C)(C)C",
        "CCC(CCC(C))C",
        "C(C(C))C"  # 二阶括号测试用例
    ]})
# 测试一阶索引
df['Indexed SMILES (Depth 1)'] = df['SMILES'].apply(lambda x: smile_grouping(x, target_depth=1))
# 测试二阶索引
df['Indexed SMILES (Depth 2)'] = df['SMILES'].apply(lambda x: smile_grouping(x, target_depth=2))
print(df)

关键修改说明 #

1. 修复分支计数错误：

   • 新增branch_last_counts列表记录每个目标层级分支的最后计数，确保新的独立分支从之前分支的最后计数+1开始，而非继承嵌套内的计数。
   • 调整括号闭合时的状态处理，正确重置当前目标计数。

2. 支持指定括号阶数：

   • 新增参数target_depth，默认值为1（原逻辑），可指定任意正整数阶数。
   • 调整层级判断逻辑，仅对open_brackets == target_depth的括号内原子进行顺序计数，更深层级的嵌套原子继承当前目标分支的计数，其余原子标记为0。

3. 鲁棒性优化：

   • 处理RDKit解析原子数与SMILES字符数不一致的场景，确保结果长度匹配。

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