Longest Univalue Path

# Definition for a binary tree node.
# class TreeNode(object):
#     def __init__(self, x):
#         self.val = x
#         self.left = None
#         self.right = None

class Solution(object):

    def longestUnivaluePath(self, root):
        """
        :type root: TreeNode
        :rtype: int
        """
        if root == None:
            return 0

        cnt = 0
        #self.max_path = 0

        max_path = [0]

        def dfs(root):
            if root == None:
                return 0

            left = dfs(root.left)
            right = dfs(root.right)
            tmp1 = 0
            tmp2 = 0

            if root.left and root.left.val == root.val:
                tmp1 = left + 1
            if root.right and root.right.val == root.val:
                tmp2 = right + 1

            if (tmp1 + tmp2) > max_path[0]:
                max_path[0] = tmp1 + tmp2

            return max(tmp1, tmp2) 


        dfs(root)
        return max_path[0]

If input data is presented by list, for example:
Vertex, V=[2,2,3], means there are three vertex v1, v2, v3 with weighted value 2, 2, 3 respectively
Edges, E=[1,2,1,3], where edge1 is (v1, v2), edge2 is (v1,v3)

#!/usr/bin/env python
from Queue import PriorityQueue

class TreeNode(object):
    def __init__(self, x):
        self.val = x
        self.neighbor = set()

if __name__ == '__main__':

    # V = [4,4,4,2,2,2]
    # E = [1,2,1,3,2,4,2,5,2,6]
    V = [1,1,3,1,3,1]
    E = [1,1,2,3,2,4,2,5,4,6]

    tree = {}

    for i in range(len(E)/2):
        idx1 = E[i*2]
        idx2 = E[i*2-1]

        if not idx1 in tree:
            node1 = TreeNode(V[idx1-1])
            tree[idx1] = node1
        else:
            node1 = tree[idx1]

        if not idx2 in tree:
            node2 = TreeNode(V[idx2-1])
            tree[idx2] = node2
        else:
            node2 = tree[idx2]

        node1.neighbor.add(node2)
        node2.neighbor.add(node1)

    visited = set()
    max_len = 0

    def dfs(node, parent_val):
        global max_len

        if node in visited:
            return 0

        visited.add(node)

        q = PriorityQueue()
        for v in node.neighbor:
            q.put(-dfs(v, node.val))

        max1 = 0
        max2 = 0
        if q.qsize() >= 2:
            max1 = -q.get()
            max2 = -q.get()
            max_len = max(max_len, max1+max2)
        elif q.qsize == 1:
            max1 = -q.get()
            max_len = max(max_len, max1)

        if parent_val == node.val:
            return max1 + 1
        else:
            return 0


    dfs(tree[1], -V[0])
    return max_len