Drunken walker emulator for EDAN55

drunk.py

#!/usr/bin/env python
# encoding: utf-8
import networkx as nx
import sys, random

def parse(file):
    f = open(file, 'r')
    G = nx.DiGraph()
    first = True
    for row in f: 
        if first:
            first = False
            G.add_nodes_from(range(int(row)), visit=0)
        else:
            nrs = row.split()
            i = 0
            while i < len(nrs) - 1:
                G.add_edge(int(nrs[i]), int(nrs[i+1]))
                i += 2
    return G
    
def drunk_walk(graph, iterations):

    damping_factor = 0.15
    current_node = 0
    for i in xrange(iterations):
        graph.node[current_node]['visit'] += 1
        if random.random() < damping_factor or len(graph.out_edges(current_node)) == 0 :
            current_node = random.choice(graph.nodes())
            #print current_node, " random jump"
        else:
            current_node = random.choice(graph.out_edges(current_node))[1]
            #print current_node

    n = len(graph.nodes())
    for i in range(n):
        print str(i) + " : " + str(graph.node[i]['visit'] / float(iterations)) 

if __name__ == "__main__":
    if len(sys.argv) == 1: 
        print "Call like: python drunk.py <nbr of iterations>"
        exit()
    
    iterations = int(sys.argv[1])
    if len(sys.argv) == 2:
        graph = parse('data/three.txt')
    else:
        graph = parse('data/' + sys.argv[2] +".txt")
    drunk_walk(graph, iterations)