Combining list of lists based on value in python

Question:

The topic was raised due to dealing with some graph problems. So I’m receiving information about graph structure via stdin and input looks like this:

Number of Edges
ID of NodeN1  Id of NodeN2
ID of NodeN3  Id of NodeN4
ID of NodeN5  Id of NodeNn
...

The first line represents amount of edges in graph, the other lines represent edges themselves(displayed as two IDs of corresponding nodes of that particular edge), for example the graph with 3 edges and 3 nodes will look like this:

3
0 1
0 2
1 2

After receiving this data, I’m parsing it in a following way, iterating over the range of input and then saving nodes in two different lists using map() approach:

E = int(input())
lst = []
lst_ = []

for x in range(E):
    l, q = map(int, input().split())
    lst.append(l)
    lst_.append(q)

After this I can just simply use something like Networkx add these edges to Graph() class, for example in the following manner:

import networkx as nx
G = nx.Graph()
for x in range (len(lst)):
    G.add_edge(lst[x],lst_[x])

Right now my problem is following – I would like to change structure of graph representation in this way: to have list of lists, where lists inside represents the nodes connected to particular node. It might sounds a bit fuzzy so let’s use example. For the graph above(with 3 nodes and 3 edges) the new list of lists structure should look like this:

mega_lst = [[1,2],[0,2],[0,1]]

First list looks like [1,2] because these are the IDs that were connected to 0 in the graph, etc. I can think of some solution using pandas indexing but I think there should be much easier way to solve it using nested lists. Can you please give a hint in what direction should I go?

Many thanks 🙂

Asked By: Keithx

||

Answers:

There may be built-in methods for this in the networkx package, but you can also do it manually like this:

inputList = [3, [0, 1], [0, 2], [1, 2]]

numNodes = inputList[0]
edges = inputList[1:]
mega_lst = []
for node in range(numNodes):
    connectedNodes = []
    for edge in edges:
        if node in edge:
            connectedNode = edge[0] if edge[1] == node else edge[1]
            connectedNodes.append(connectedNode)
    mega_lst.append(connectedNodes)

print(mega_lst)
Answered By: Temba
Categories: questions Tags: ,
Answers are sorted by their score. The answer accepted by the question owner as the best is marked with
at the top-right corner.