二度人脉计算

编程实现二度人脉的计算

__author__ = '20130907'
import networkx,heapq,sys
from matplotlib import pyplot
from collections import defaultdict,OrderedDict
from numpy  import array

def Edge():
    return defaultdict(Edge)

class Graph:
	def __init__(self):
		self.Link = Edge()
		self.FileName = ''
		self.Separator = ''

	def MakeLink(self,filename,separator):
		self.FileName = filename
		self.Separator = separator
		graphfile = open(filename,'r')
		for line in graphfile:
			items = line.split(separator)
			self.Link[items[0]][items[1]] = int(1)
			self.Link[items[1]][items[0]] = int(1)
		graphfile.close()
	def NDegreeNode(self,start,n):
		pathFromStart = {}
		pathFromStart[start] = [start]
		pathLenFromStart = {}
		pathLenFromStart[start] = 0
		todoList = [start]
		while todoList:
			current = todoList.pop(0)
			for neighbor in self.Link[current]:
				if neighbor not in pathFromStart:
					pathFromStart[neighbor] = pathFromStart[current] + [neighbor]
					pathLenFromStart[neighbor] = pathLenFromStart[current] + 1
					if pathLenFromStart[neighbor] <= n+1:
						todoList.append(neighbor)

		for node in pathFromStart.keys():
			if len(pathFromStart[node]) != n+1:
				del pathFromStart[node]
		return pathFromStart
	def Draw(self):
		G = networkx.Graph()
		nodes = self.Link.keys()
		edges = [(node,neighbor) for node in nodes for neighbor in self.Link[node]]
		G.add_edges_from(edges)
		networkx.draw(G)
		pyplot.show()

if __name__=='__main__':
	separator = '\t'
	filename = r'C:\Users\20130907\Desktop\uid.txt'
	resultfilename = 'C:\Users\20130907\Desktop\result.txt'

	myGraph = Graph()
	myGraph.MakeLink(filename,separator)
	print 'NDegreeNode',myGraph.NDegreeNode('111360',2).keys()
	myGraph.Draw()