networkx: contract adjacent nodes in a directed acyclic graph

Question

Zoom into my directed acyclic graph , you will find the adjacent nodes (u,v,w,x,y) . The node w is a the center, u and v are upstream and x and y are downstream

That results in the following edges:

(u,w)
(v,w)
(w,x)
(w,y)

I want to remove the node w and preserve the flow of the graph. After w removed, the edges should be

(u,x)
(u,y)
(v,x)
(v,y)

I thought that networkx.contracted_nodes(G, u, w, self_loops=False) is the way to go. But according the docs this will generate a new node, let's call it uw with the edges

(uw,x)
(uw,y)
(v,uw)

Which alters the behaviour of my graph and is not what I want.

Is there a way to solve it using.networkx?

Answer 1

Assuming edges are a list of tuples

def foo(arr, key):
    d = {}
    for a, b in arr:
        d.setdefault(a, []).append(b)
    
    if d.get(key) is None:
        return arr
    
    ans = []
    fill_vals = d[key]
    for a, b in arr:
        if a != key:
            if b != key:
                ans.append((a, b))
            else:
                for val in fill_vals:
                    ans.append((a, val))
                    
    return ans

data = [("u", "w"), ("v", "w"), ("w", "x"), ("w", "y")]    
foo(data, "w")
# [('u', 'x'), ('u', 'y'), ('v', 'x'), ('v', 'y')]