<p>In this problem, a rooted tree is a <b>directed</b> graph such that, there is exactly one node (the root) for which all other nodes are descendants of this node, plus every node has exactly one parent, except for the root node which has no parents.</p>

<p>The given input is a directed graph that started as a rooted tree with <code>n</code> nodes (with distinct values from <code>1</code> to <code>n</code>), with one additional directed edge added. The added edge has two different vertices chosen from <code>1</code> to <code>n</code>, and was not an edge that already existed.</p>

<p>The resulting graph is given as a 2D-array of <code>edges</code>. Each element of <code>edges</code> is a pair <code>[u<sub>i</sub>, v<sub>i</sub>]</code> that represents a <b>directed</b> edge connecting nodes <code>u<sub>i</sub></code> and <code>v<sub>i</sub></code>, where <code>u<sub>i</sub></code> is a parent of child <code>v<sub>i</sub></code>.</p>

<p>Return <em>an edge that can be removed so that the resulting graph is a rooted tree of</em> <code>n</code> <em>nodes</em>. If there are multiple answers, return the answer that occurs last in the given 2D-array.</p>

<p>&nbsp;</p>
<p><strong class="example">Example 1:</strong></p>
<img alt="" src="https://assets.leetcode.com/uploads/2020/12/20/graph1.jpg" style="width: 222px; height: 222px;" />
<pre>
<strong>Input:</strong> edges = [[1,2],[1,3],[2,3]]
<strong>Output:</strong> [2,3]
</pre>

<p><strong class="example">Example 2:</strong></p>
<img alt="" src="https://assets.leetcode.com/uploads/2020/12/20/graph2.jpg" style="width: 222px; height: 382px;" />
<pre>
<strong>Input:</strong> edges = [[1,2],[2,3],[3,4],[4,1],[1,5]]
<strong>Output:</strong> [4,1]
</pre>

<p>&nbsp;</p>
<p><strong>Constraints:</strong></p>

<ul>
	<li><code>n == edges.length</code></li>
	<li><code>3 &lt;= n &lt;= 1000</code></li>
	<li><code>edges[i].length == 2</code></li>
	<li><code>1 &lt;= u<sub>i</sub>, v<sub>i</sub> &lt;= n</code></li>
	<li><code>u<sub>i</sub> != v<sub>i</sub></code></li>
</ul>

import heapq
# https://leetcode.com/problems/redundant-connection/
# really good dj set medium! Lots of graph qs probably can have a disjoint set soln.
# hard

# burte force worked! remove each edge, see if its connected.
# maybe need to do that with a disjoint set? doesn't look like we can do that...
class DisjointSet: 
    
    def __init__(self, n):
        self.disjoint_set = [-1 for _ in range(n + 1)]

    def find(self, i): # find the representation of i
        group_index = self.disjoint_set[i]
        if group_index < 0:
            return i
        
        group_index = self.find(group_index)
        # path compression, join this set to the root.
        self.disjoint_set[i] = group_index

        return group_index

    def join(self, a_group_index, b_group_index):
        # assume they are not the same.
        
        # root is negative size, smaller = larger
        size_a = self.disjoint_set[a_group_index]
        size_b = self.disjoint_set[b_group_index]

        if size_a < size_b: # a is bigger, join smaller to bigger.
            self.disjoint_set[a_group_index] += self.disjoint_set[b_group_index]
            self.disjoint_set[b_group_index] = a_group_index
        else:
            # b is bigger, join smaller to bigger
            self.disjoint_set[b_group_index] += self.disjoint_set[a_group_index]
            self.disjoint_set[a_group_index] = b_group_index    


class Solution:
    def dfs(self, node, visited, graph): 
        visited.add(node)
        for adj in graph[node]['edges']:
            if adj in visited: # it shouldn't be, causes a cycle, we don't have our ans
                return
            self.dfs(adj, visited, graph)
        

    def findRedundantDirectedConnection(self, edges: List[List[int]]) -> List[int]:
        # build graph
        # remove edge
        # check if connected
        if not edges: return None 

        num_nodes = 0
        graph = {} # { node_v: { indeg: 0, edges: { node_c: edge_index_weight }}}
        # MST problem from a topolgical ordering / find the root and build it up that way.
        for i, e in enumerate(edges):
            if e[0] not in graph:
                graph[e[0]] = {'indeg': 0, 'edges': {}}
            if e[1] not in graph:
                graph[e[1]] = {'indeg': 0, 'edges': {}}
            graph[e[0]]['edges'][e[1]] = i
            graph[e[1]]['indeg'] += 1
            num_nodes = max(e[0], e[1], num_nodes)
        
        for i in range(len(edges) -1, -1, -1):
            edge = edges[i]
            # try remove this edge.
            del graph[edge[0]]['edges'][edge[1]]
            # remove edge[0] -> edge[1]
            graph[edge[1]]['indeg'] -= 1

            # find root.
            root = None 
            for start in graph:
                if graph[start]['indeg'] == 0:
                    root = start # todo might have many roots handle later.
                    break
            
            if root:
                # start traversal from root.
                visited = set()
                self.dfs(root, visited, graph)
                # if graph connected we are ok.
                # nodes are lablled 1-n so # of nodes in the connected graph is n.
                if len(visited) == num_nodes:
                    # print(visited, root, num_nodes)
                    # we have an answer.
                    return edge
            # add the edge back.
            graph[edge[0]]['edges'][edge[1]] = i
            graph[edge[1]]['indeg'] += 1



        # if edges == [[1,2],[3,1],[2,3]]:
        #     return [2,3]

        n = len(edges)
        dj = DisjointSet(n)
        # close but no since we don't know the root and it might have an in deg.
        # maybe this is cycle detection? 
        if not edges: 
            return None 

        graph = {} # { node_v: { indeg: 0, edges: { node_c: edge_index_weight }}}
        # MST problem from a topolgical ordering / find the root and build it up that way.
        for i, e in enumerate(edges):
            if e[0] not in graph:
                graph[e[0]] = {'indeg': 0, 'edges': {}}
            if e[1] not in graph:
                graph[e[1]] = {'indeg': 0, 'edges': {}}
            graph[e[0]]['edges'][e[1]] = i
            graph[e[1]]['indeg'] += 1

        # there is a root that exists, with 0 indegree
        # there is exactly one node (the root) for which all other nodes are descendants of this node
        root = None
        for node in graph:
            if graph[node]['indeg'] == 0:
                root = node
        
        if not root: 
            # pretend root is the first edge.
            # might err here? 
            root = edges[0][0]
       
        # and weighted edges on the index.
        reachable = []
        for edge in graph[root]['edges']:
            w = graph[root]['edges'][edge]
            e = (root, edge)
            reachable.append((w, e))
        heapq.heapify(reachable)

        seen_nodes = set([root]) 

        while reachable:
            min_edge = heapq.heappop(reachable)
            new_node = min_edge[1][1]
            if new_node in seen_nodes:
                # each node has 1 parent. 
                # i know new_node has been seen before so we have accoutned for (p -> new_node)
                # we also get (min_edge[0] -> new_node)
                # search edges for which comes last.
                
                def get_next_edge(edge):
                    if not edge:
                        return None 
                    if not edge[1] in graph:
                        return None 
                    connected_edges = graph[edge[1]]['edges']
                    # asusme there is only 1? 
                    if not connected_edges:
                        return None
                    for x in connected_edges: # hacky
                        return (edge[1], x)

                # if cycle, will effect the tree, remove what is NOT on the path from the root?
                has_cycle = False
                slow = fast = min_edge[1]

                while slow and fast:
                    print(slow, fast)
                    slow = get_next_edge(slow)
                    fast = get_next_edge(fast)
                    fast = get_next_edge(fast)

                    print(slow, fast)
                    if slow == None:
                        break
                    if fast == None:
                        break
                    if slow == fast:
                        has_cycle = True
                        break 
                # if no cycle, doesn't effect the tree, remove last
                for i in range(len(edges) - 1, -1, -1):
                    if has_cycle:
                        print('aa')
                        # if edges[i][1] == new_node and root != edges[i][0]:
                        #     return edges[i]
                        return min_edge[1] # return not on the root path 
                    else:
                        print('bb')
                        if edges[i][1] == new_node:
                            return edges[i]
                print('cc')
            else:
                seen_nodes.add(new_node)
                # add edges.
                for edge in graph[new_node]['edges']:
                    w = graph[new_node]['edges'][edge]
                    e = (new_node, edge)
                    heapq.heappush(reachable, (w, e))
                    

        # start at root, BFS take min from heap.
        


        # for src, dest in edges: 
        #     src_group = dj.find(src)
        #     dest_group = dj.find(dest)
        #     if src_group == dest_group:
        #         return [src, dest]
        #     dj.join(src_group, dest_group)

        
17700