import networkx as nx
import matplotlib.pyplot as plt

# read in our graph, as a list of edges
G = nx.read_edgelist("lec01.data")

G.order()   # number of vertices
G.size()    # number of edges
G.nodes     # list of vertices
G.edges     # list of edges
G.adj       # adjacencies of each vertex

# draw and display our graph
nx.draw(G, with_labels=True)
plt.show()

# run our connectivity algorithm
counter = 0
S = {}
for v in G.nodes:
    S[v] = counter
    counter += 1

updates = 1
while updates > 0:
    updates = 0
    for v in G.nodes:
        for u in G.neighbors(v):
            if S[u] > S[v]:
                S[v] = S[u]
                updates += 1

# bit of a hack for visualization - won't generalize to other graphs
colors = []
for v in G.nodes:
    if S[v] == 7:
        colors.append('green')
    elif S[v] == 4:
        colors.append('red')

nx.draw(G, with_labels=True, node_color=colors)
plt.show()