python - 返回具有相同值节点的最长路径

Question

我刚刚在研究一些 Google 面试问题，我在网上偶然发现了这个我似乎无法理解的问题

考虑一棵具有 N 个节点的无向​​树，编号从 1 到 N。每个节点都有一个与之关联的标签，该标签是一个整数值。不同的节点可以有相同的标签。编写一个函数，给定一个长度为 N 的零索引数组 A，其中 A[j] 是树中第 (j + 1) 个节点的标签值，以及一个长度为 K = (N – 1) * 2 其中描述了树的边，返回最长路径的长度，使得该路径上的所有节点都具有相同的标签。长度是该路径中的边数。

例子:

A = [1, 1, 1, 2, 2]E = [1, 2, 1, 3, 2, 4, 2, 5]

这棵树如下所示。节点遵循表单标签、值。

----------1, 1

-----1, 2        1, 3

2, 4      2, 5

函数应该返回2，因为最长的路径是2->1->3，这条路径有2条边。

假设 1 <= N <= 1000 并且数组 A 的每个元素都是 [1, 1000000000] 范围内的整数。

您对打破这个解决方案有何看法？谢谢!

Answer 1

这是我的建议。我没有仔细检查过。

from collections import defaultdict

A = [1, 1, 1, 2, 2] 

E = [1, 2, 1, 3, 2, 4, 2, 5]

d = defaultdict(list)
it = iter(E)
for k, v in zip(it, it):
    d[k].append(v)
    d[v].append(k)

leaves = [k for k, v in d.items() if len(v) == 1]

def len_path(node, length=0, coming_from=None):

    max_loc_len = length
    loc_data = A[node-1]

    available_nodes = {i: A[i-1] for i in d[node]}
    available_nodes.pop(coming_from, None)

    for i, i_data in available_nodes.items():
        cumul = length + 1 if loc_data == i_data else 0
        loc_len = len_path(i, cumul, node)
        if loc_len > max_loc_len:
            max_loc_len = loc_len 

    return max_loc_len

max_len = max([len_path(leaf) for leaf in leaves])

以及更复杂的搜索，避免两次计算相同的路径:

# ...
def len_path(node, length=0, coming_from=None, fork=False):

    max_loc_len = length
    loc_data = A[node-1]

    available_nodes = {i: A[i-1] for i in d[node]}
    available_nodes.pop(coming_from, None)

    matching_nodes = [k for k, v in available_nodes.items()
                         if v == loc_data and k not in leaves[:explored]]
    if len(matching_nodes) > 1:
        fork = True

    if not available_nodes and not fork and node in leaves[explored:]:
        # Reached an unexplored leaf without forks on the path, 
        # hence no need to explore it later
        leaves.remove(node)

    for i, i_data in available_nodes.items():
        cumul = length + 1 if loc_data == i_data else 0
        loc_len = len_path(i, cumul, node, fork)
        if loc_len > max_loc_len:
            max_loc_len = loc_len 

    return max_loc_len

explored = 0
max_len =0

while explored < len(leaves):
    length = len_path(leaves[explored])
    if length > max_len:
        max_len = length
    explored += 1