查找树内最近的节点

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