查找树内最近的节点

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

例如，Katy 将与 Tom 匹配。由于与 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