python - 寻找树内最近的节点

Question

我有一个 pandas.DataFrame 包含树中的节点。该表如下所示:

╔═══════╦════════╦════════╦══════╗
║ index ║ color  ║  name  ║ head ║
╠═══════╬════════╬════════╬══════╣
║     0 ║ red    ║ Tom    ║    0 ║
║     1 ║ blue   ║ Lucy   ║    0 ║
║     2 ║ green  ║ Peter  ║    1 ║
║     3 ║ red    ║ Katy   ║    1 ║
║     4 ║ green  ║ Sam    ║    4 ║
║     5 ║ orange ║ Linda  ║    2 ║
║     6 ║ blue   ║ Robert ║    4 ║
║     7 ║ brown  ║ James  ║    6 ║
║     8 ║ red    ║ Betty  ║    7 ║
║     9 ║ red    ║ Amanda ║    4 ║
║    10 ║ black  ║ Luke   ║    8 ║
╚═══════╩════════╩════════╩══════╝

列head存储父节点的索引。它将创建一棵树，如下所示:

并且每个节点可以有0+个子节点(不限于2个)。

当我选择一个人时，我想找到另一个具有相同颜色的人。有 3 条规则:

1. 如果在同一个词干上，则选择最近的人
2. 如果没有选择任何人，请选择同一棵树中最近的人
3. 如果无法选择任何人，则返回None

例如，凯蒂将与汤姆匹配。由于与 Betty 的同一茎不再有红色，因此将选择 Amanda。

除了暴力破解之外，还有什么方法可以得到答案吗？

Answer 1

我使用了网络分析技术，不确定它是否最适合您的情况。

这个想法很简单:

1. 制作网络图
2. 找到与您所选人员颜色相同的所有其他人，我将其称为候选人
3. 检查候选人和所选人员在网络中是否连接(即候选人和所选人员之间是否存在路径)
4. 找到最短路径的候选者

这是我的代码

import io
import pandas as pd
import networkx as nx
from networkx.algorithms import shortest_path, has_path


# Data
df_str = """
index,colour,name,head
0,red,Tom,0
1,blue,Lucy,0
2,green,Peter,1
3,red,Katy,1
4,green,Sam,4
5,orange,Linda,2
6,blue,Robert,4
7,brown,James,6
8,red,Betty,7
9,red,Amanda,4
10,black,Luke,8
"""
df = pd.read_csv(io.StringIO(df_str), sep=",")


# Function to find the closest person with the same colour as the person with `id0`
def find_same_colour(id0, df):
    # Create network
    g = nx.Graph()
    for _, row in df.iterrows():
        g.add_node(row['index'], colour=row['colour'])
        if row['index'] != row['head']:
            g.add_edge(row['index'], row['head'])
    # List out candidates
    colour = df.loc[df['index'].values == id0, 'colour'].values[0]
    candidates = df.loc[(df['colour'].values == colour) & (df['index'].values != id0), 'index'].values
    # Initialise searching
    target = None
    l_max = df.shape[0] + 2
    # Search
    for i in candidates:
        if has_path(g, id0, i):
            path = shortest_path(g, id0, i)
            if len(path) < l_max:
                target = i
    return target


for i in df['index'].values:
    print(i, find_same_colour(i, df), sep='-')

这是输出，

# 0-3
# 1-None
# 2-None
# 3-0
# 4-None
# 5-None
# 6-None
# 7-None
# 8-9
# 9-8
# 10-None