endpoints of tree diameter

'''
given you count of treeNodes and the edges
mark the endpoint nodes, which are part of diameter

'''

class Solution:

    def calculateddp(self,ddp,root,treeMap):
        if root not in treeMap:
            ddp[root] = 1
            return
        maxDown = 0
        for c in treeMap[root]:
            self.calculateddp(ddp,c,treeMap)
            maxDown = max(maxDown,ddp[c])
        ddp[root] = 1 + maxDown
        return


    def endPointsOfTreeDiameter(self,nodes,nodeFrom,nodeTo):
        treeMap = {}
        parentMap = {}
        edgesLen = len(nodeFrom)
        for i in range(edgesLen):
            if nodeFrom[i] in treeMap:
                treeMap[nodeFrom[i]].append(nodeTo[i])
            else:
                treeMap[nodeFrom[i]] = [nodeTo[i]]
            parentMap[nodeTo[i]] = nodeFrom[i]

        root = 1 #assuming root to be 1
        for i in range(1,nodes+1):
            if i not in nodeTo:
                root = i
                break
        #calculate downward maxlen dp for every node
        ddp = [0 for _ in range(nodes+1)]
        self.calculateddp(ddp,root,treeMap)

        #calculate upward dp as max(curr + parent + sibling downward dp's) #curr is 1 , parent is 1 if exist
        udp = [0 for _ in range(nodes+1)]

        parent = None
        q = []
        q.extend(treeMap[root]) #you cannot attach treeMap[parent] otherwise
        udp[root] = 1
        while len(q):
            curr = q.pop(0)
            parent = parentMap[curr]#getting the parent for curr
            # print(f"treeMap[parent] {treeMap[parent]}")
            # print(f"parent is {parent} treeMap[parent] is {treeMap[parent]}")
            siblings = [x for x in treeMap[parent] if x!=curr]
            # print(f"curr is {curr} siblings {siblings}")
            siblingsddp = [ddp[x] for x in siblings]
            # print(f"curr is {curr} siblingsddp is {siblingsddp}")
            maxsiblingddp = max(siblingsddp) if len(siblingsddp) else 0
            udp[curr] = 1+max(1+maxsiblingddp,udp[parent])#(1stroke up and then downward direction) vs (udp of parent)
            if len(treeMap.get(curr,[])):
                q.extend(treeMap[curr])#imp

        # maxi = max(udp) #getting the diameter
        # print(f"ddp is {ddp}")

        # print(f"udp is {udp}")
        # print(f"maxi is {maxi}")
        maxdiams = [0 for _ in range(nodes+1)]
        ans = [0 for _ in range(nodes+1)]
        for i in range(len(maxdiams)):
            maxdiams[i] = max(udp[i],ddp[i])#either one length can you give you max path ,so even if root is endpoint we can handle
        maxi = max(maxdiams)
        for i in range(len(maxdiams)):
            if maxi == maxdiams[i]:
                ans[i] = 1


        if udp[root]+ddp[root]-1 == maxi:
            ans[root] = 1
        return ans[1:]



def run_tests():
    sol = Solution()
    cases = [
        (7, [1,2,3,3,1,1], [2,3,4,5,6,7]),
        (9, [1,2,8,3,3,1,1,8], [2,8,3,4,5,6,7,9]),
        (2, [2], [1]),
        (5, [1,1,2,2], [2,3,4,5]),
        (3, [1,2], [2,3]),
        (4, [1,2,2], [2,3,4])
    ]

    for case in cases:
        print(sol.endPointsOfTreeDiameter(*case))


if __name__ == "__main__":
    run_tests()