Asignacion de gastos en publicidad oficial de un organismo a un proveedor por rubro en base a vistas agregadas

main.py

import csv
import networkx as nx
from decimal import *
from itertools import *

FILE_SALIDA = 'output.csv'
FILE_PROVEEDOR = 'ocrtest_segundo_semestre_prov_rubro CORREGIDO.csv' 
FILE_ORGANISMO = 'ocrtest_segundo_semestre_organismo_rubro CORREGIDO.csv' 

class Data:
	def __init__(self):
		self.proveedores = []
		self.organismos = []
		self.rubros = ["cable", "cine", "grafica", "radio", "tv", "via publica", "web"]
		self.gasto_proveedor = {}
		self.gasto_organismo = {}
		for rubro in self.rubros:
			self.gasto_proveedor[rubro] = []
			self.gasto_organismo[rubro] = []
		
	def fill_from_csv(self, file_proveedor, file_organismo):
		with open(file_proveedor, 'rb') as f:
			prov_reader = csv.reader(f, delimiter=',', quotechar='"')
			for row in prov_reader:
				self.proveedores.append(row[0])
				self._add_gasto(row, self.gasto_proveedor,1)
		
		with open(file_organismo, 'rb') as f:
			org_reader = csv.reader(f, delimiter=',', quotechar='"')
			for row in org_reader:
				self.organismos.append(row[0])
				self._add_gasto(row, self.gasto_organismo,2)
		return self
				
	def _add_gasto(self, row, dict, offset):
		name = row[0]
		for i in range(7):
			gasto = row[i+offset]
			if gasto: gasto = Decimal(gasto.replace(',','').replace('$',''))
			else: gasto = Decimal('0')
			dict[self.rubros[i]].append((name, gasto))

class Solver:
	def __init__(self, data):
		self.data = data
			
	def solve(self, rubro=None):
		if rubro: 
			return self._solve_rubro(rubro)
		solution = {}
		for rubro in self.data.rubros:
			for org,prov,gasto in self._solve_rubro(rubro):
				solution[(org,prov,rubro)] = gasto
		return solution
		
	
	def _solve_rubro(self,rubro):
		graph = nx.DiGraph(name="Publicidad Oficial para " + rubro)
		graph.add_edge('dummy','source')
		for (org,gasto) in self.data.gasto_organismo[rubro]:
			graph.add_edge('source', org, capacity = gasto)
		for (prov,gasto) in self.data.gasto_proveedor[rubro]:
			graph.add_edge(prov, 'sink', capacity = gasto)
		for (org,prov) in product(self.data.organismos, self.data.proveedores):
			graph.add_edge(org,prov)
		flow = nx.ford_fulkerson_flow(graph, 'source', 'sink')
		return [(v,w,flow[v][w]) for v,w in graph.edges() if not v == 'source' and not w == 'sink' and not w == 'source']

def test():
	data = Data()
	data.proveedores = ['P1', 'P2']
	data.organismos = ['O1', 'O2', 'O3']
	data.rubros = ['R1']
	data.gasto_proveedor = {'R1': [('P1', 100.0), ('P2', 400.0)]}
	data.gasto_organismo = {'R1': [('O1', 200.0), ('O2', 300.0), ('O3', 0.0)]}
	solver = Solver(data)
	for (org,prov,gasto) in solver.solve(data.rubros[0]):
		print org, prov, gasto
	
def main():
	data = Data().fill_from_csv(FILE_PROVEEDOR, FILE_ORGANISMO)
	solver = Solver(data)
	gastos = solver.solve()
	
	with open(FILE_SALIDA, 'wb') as f:
		writer = csv.writer(f, delimiter=',', quotechar='"', quoting=csv.QUOTE_NONNUMERIC)
		writer.writerow(['Organismo', 'Proveedor'] + data.rubros)
		for (org,prov) in product(data.organismos, data.proveedores):
			writer.writerow([org,prov] + [gastos[(org,prov,rubro)] for rubro in data.rubros])

	
if __name__ == "__main__":
    main()