def bestMap(L1, L2): import numpy as np from hungarian import hungarian L1 = np.array(L1) L2 = np.array(L2) assert (L1.size == L2.size) L1 = L1.astype(np.int32) L2 = L2.astype(np.int32) Label1 = list(set(L1)) nClass1 = len(Label1) Label2 = list(set(L2)) nClass2 = len(Label2) nClass = max(nClass1, nClass2) G = np.zeros((nClass, nClass), dtype=np.int32) for i in range(nClass1): for j in range(0, nClass2): G[i, j] = len(np.where((L1 == Label1[i]) & (L2 == Label2[j]))[0]) c, t = hungarian(-G) newL2 = np.zeros(L2.size, dtype=np.int32) for i in range(nClass2): newL2[L2 == Label2[i]] = Label1[c[i] - 1] return newL2
def matchFlies(self, cost, maxcost): # number of targets from previous frame ntargets = cost.shape[1] # number of observations in current frame nobs = cost.shape[0] # try greed assignment obsfortarget = cost.argmin(axis=0) # make sure not assigning obs when should be assigning to dummy mincost = cost.min(axis=0) obsfortarget[mincost > maxcost] = -1 isconflict = 0 isunassigned = np.empty((nobs, 1)) isunassigned[:] = True for i in range(ntargets): if obsfortarget[i] < 0: continue if isunassigned[obsfortarget[i]] == False: isconflict = True isunassigned[obsfortarget[i]] = False if not isconflict: print "Using brute force" return (obsfortarget, isunassigned) # otherwise, use hungarian matching algorithm n = ntargets + nobs weights = np.zeros((n,n)) weights[0:nobs, 0:ntargets] = cost weights[0:nobs, ntargets:n] = maxcost weights[nobs:n, 0:ntargets] = maxcost (targetforobs, obsfortarget) = hungarian(weights) # remove dummy targets obsfortarget = obsfortarget[0:ntargets] # unassign obs assigned to dummy targets isunassigned = targetforobs[0:nobs] >= ntargets obsfortarget[obsfortarget >= nobs] = -1 print "Using hungarian" return (obsfortarget, isunassigned)
def matchidentities( cost, maxcost=-1, issplit=None, maxcost_split=-1 ): """(observationfortarget,unassignedobservations) = matchidentities( cost,maxcost ) An observation is a new piece of information, i.e., something we'd like to correlate with what we've previously seen. A target is an old piece of information, e.g., a position where we might reasonably expect an observation to appear. 'cost' is a n_observations x n_targets matrix, where cost[i,j] is the cost of assigning observation i to target j 'maxcost' is the maximum distance between a target and its assigned observation 'observationfortarget' is a n_targets length array, where observationfortarget[i] is the index of the observation assigned to target i 'isunassignedobservation' is a n_observations length vector, where isunnassignedobservation[i] is True if the observation is not assigned to any target.""" # KB 20120109: allow jumps of distance max_jump_split if an observation is the result of # splitting a connected component # TODO: this raises errors when n_observations and(/or?) n_targets == 0 if maxcost < 0: maxcost = params.max_jump # KB 20120109: if maxcost_split not input, then use max_jump_split if maxcost_split < 0: maxcost_split = params.max_jump_split # KB 20120109: do we need to deal with a different cost for splits handlesplits = issplit is not None and \ maxcost != maxcost_split and \ issplit.any() # number of targets in the previous frame ntargets = cost.shape[1] # number of observations in the current frame nobservations = cost.shape[0] # try greedy: assign each target its closest observation observationfortarget = cost.argmin(axis=0) # make sure we're not assigning observations when we should # be assigning a lost target mincost = cost.min(axis=0) # KB 20120109: only use the greedy method if we're not using a different cost for splits if not handlesplits: observationfortarget[mincost>maxcost] = -1 # see if there are any conflicts: count the number of targets claiming # each observation isconflict = 0 # initialize whether an observation is unassigned to a target to be all True isunassignedobservation = num.empty((nobservations,1)) isunassignedobservation[:] = True for i in range(ntargets): if observationfortarget[i] < 0: continue if isunassignedobservation[observationfortarget[i]] == False: isconflict = True isunassignedobservation[observationfortarget[i]] = False if isconflict == False: # greedy is okay, so just return return ( observationfortarget, isunassignedobservation ) # if there is a conflict, then use the Hungarian algorithm # create a cost matrix that is (nnodes = ntargets+nobservations) x nnodes last_time = time.time() nnodes = ntargets + nobservations costpad = num.zeros((nnodes,nnodes)) # top left square is the original cost matrix costpad[0:nobservations,0:ntargets] = cost # top right square is maxcost costpad[0:nobservations,ntargets:nnodes] = maxcost # KB 20120109: if any observations are split, they get a different max cost if handlesplits: costpad[issplit,ntargets:nnodes] = maxcost_split # bottom left square is maxcost costpad[nobservations:nnodes,0:ntargets] = maxcost # optimize using hungarian method (targetforobservation,observationfortarget) = hungarian( costpad ) # we don't care about the dummy target nodes observationfortarget = observationfortarget[0:ntargets] # observations assigned to dummy target nodes are unassignedd isunassignedobservation = targetforobservation[0:nobservations] >= ntargets observationfortarget[observationfortarget>=nobservations] = -1 return (observationfortarget,isunassignedobservation)
def matchidentities( cost, maxcost=-1 ): """(observationfortarget,unassignedobservations) = matchidentities( cost,maxcost ) An observation is a new piece of information, i.e., something we'd like to correlate with what we've previously seen. A target is an old piece of information, e.g., a position where we might reasonably expect an observation to appear. 'cost' is a n_observations x n_targets matrix, where cost[i,j] is the cost of assigning observation i to target j 'maxcost' is the maximum distance between a target and its assigned observation 'observationfortarget' is a n_targets length array, where observationfortarget[i] is the index of the observation assigned to target i 'isunassignedobservation' is a n_observations length vector, where isunnassignedobservation[i] is True if the observation is not assigned to any target.""" # TODO: this raises errors when n_observations and(/or?) n_targets == 0 if maxcost < 0: maxcost = params.max_jump # number of targets in the previous frame ntargets = cost.shape[1] # number of observations in the current frame nobservations = cost.shape[0] # try greedy: assign each target its closest observation observationfortarget = cost.argmin(axis=0) # make sure we're not assigning observations when we should # be assigning a lost target mincost = cost.min(axis=0) observationfortarget[mincost>maxcost] = -1 # see if there are any conflicts: count the number of targets claiming # each observation isconflict = 0 # initialize whether an observation is unassigned to a target to be all True isunassignedobservation = num.empty((nobservations,1)) isunassignedobservation[:] = True for i in range(ntargets): if observationfortarget[i] < 0: continue if isunassignedobservation[observationfortarget[i]] == False: isconflict = True isunassignedobservation[observationfortarget[i]] = False if isconflict == False: # greedy is okay, so just return if params.print_crap: print "Greedy is okay" return ( observationfortarget, isunassignedobservation ) # if there is a conflict, then use the Hungarian algorithm # create a cost matrix that is (nnodes = ntargets+nobservations) x nnodes last_time = time.time() nnodes = ntargets + nobservations costpad = num.zeros((nnodes,nnodes)) # top left square is the original cost matrix costpad[0:nobservations,0:ntargets] = cost # top right square is maxcost costpad[0:nobservations,ntargets:nnodes] = maxcost # bottom left square is maxcost costpad[nobservations:nnodes,0:ntargets] = maxcost if params.print_crap: print 'Need to use Hungarian: time to create cost matrix: %.2f'%(time.time() - last_time) last_time = time.time() # optimize using hungarian method (targetforobservation,observationfortarget) = hungarian.hungarian(costpad) if params.print_crap: print 'Time to optimize: %.2f'%(time.time() - last_time) last_time = time.time() # we don't care about the dummy target nodes observationfortarget = observationfortarget[0:ntargets] # observations assigned to dummy target nodes are unassignedd isunassignedobservation = targetforobservation[0:nobservations] >= ntargets observationfortarget[observationfortarget>=nobservations] = -1 return (observationfortarget,isunassignedobservation)
dim = 4 matrix_dim = 16 # random and permutation num # X = np.random.random((num, 12, 10)) # X = np.random.randint(0, 99, size=[num,dim,dim]) H = np.zeros((dim, dim), dtype=int) A = [] # selected affinity matrices P = [] # selected permutation matrices # H = np.negative(H) total = 0 while (total < amount): X = np.random.randint(0, 99, size=[dim, dim]) cost, hung = hungarian(X) solutions = [] cost, solutions = allSolutions(cost, solutions) num = len(solutions) if num == 1: for j in range(dim): H[j, solutions[0][j]] = 1 A.append(X) P.append(H) total += 1 H = np.zeros((dim, dim), dtype=int) print("Number of matrices with 1 solution: %d" % len(P)) np.save('affinity', A) np.save('solutions', P)
print "trial = " + str(trial) (x, k, w) = initialize() state = random.get_state() if True: (mu0, S0, priors0, gamma0, err0) = kcluster0.gmm(x, k, weights=w, nreplicates=10) random.set_state(state) (mu, S, priors, gamma, err) = kcluster.gmm(x, k, weights=w, nreplicates=10) D = (num.tile(mu[:,0].reshape((k,1)),(1,k)) - mu0[:,0])**2 + \ (num.tile(mu[:,1].reshape((k,1)),(1,k)) - mu0[:,1])**2 (a1, a2) = hungarian.hungarian(D) if True or \ num.max(num.abs(mu - mu0[a1,:])) > .001 or \ num.max(num.abs(S-S0[:,:,a1])) > .001 or \ num.max(num.abs(priors - priors0[a1])) > .001 or \ num.max(num.abs(gamma-gamma0[:,a1])) > .001 or \ num.abs(err - err0) > .001: print "cluster approx cluster0[" + str(a1) + "]" print "max(|mu - mu0|) = " + str( num.max(num.abs(mu - mu0[a1, :]))) print "max(|S - S0|) = " + str( num.max(num.abs(S - S0[:, :, a1]))) print "max(|priors - priors0|) = " + str( num.max(num.abs(priors - priors0[a1]))) print "max(|gamma - gamma0|) = " + str( num.max(num.abs(gamma - gamma0)))
def KS(X_1, X_2): init_type = 'eig' # random, eig, ... llambda = 1.0 # step size n_iter = 100 # number of LAP iterations omegas = 2.0 n_obs = X_1.shape[0] # number of observations bases = numpy.eye(n_obs) # normalization of data l2norm = numpy.zeros((n_obs, 1)) for i in xrange(n_obs): l2norm[i] = numpy.sum(X_1[i, ]) new_X_1 = X_1 / l2norm starttime = time.clock() # kernel for grid dl = CRBFKernel() dL = dl.Dot(X_2, X_2) omega_L = 1.0 / numpy.median(dL.flatten()) kernel_L = CGaussKernel(omega_L) L = kernel_L.Dot(X_2, X_2) # should use incomplete cholesky instead ... # kernel for data tK = numpy.zeros((n_obs, n_obs)) for i in xrange(n_obs): for j in xrange(i, n_obs): temp = ((new_X_1[i, :] - new_X_1[j, :])** 2) / (new_X_1[i, :] + new_X_1[j, :]) inan = numpy.isnan(temp) temp[inan] = 0.0 tK[i, j] = tK[j, i] = numpy.sum(temp) * 0.5 mdisK = numpy.median(numpy.median(tK)) K = numpy.zeros((n_obs, n_obs)) for i in xrange(n_obs): for j in xrange(i, n_obs): temp = ((new_X_1[i, :] - new_X_1[j, :])** 2) / (new_X_1[i, :] + new_X_1[j, :]) inan = numpy.isnan(temp) temp[inan] = 0.0 K[i, j] = K[j, i] = numpy.exp(-1.0 * omegas / mdisK * numpy.sum(temp) * 0.5) #chi-square kernel stoptime = time.clock() print 'computing kernel matrices (K and L) takes %f seconds ' % (stoptime - starttime) print 'original objective function : ' H = numpy.eye(n_obs) - numpy.ones(n_obs) / n_obs print numpy.trace(numpy.dot(numpy.dot(numpy.dot(H, K), H), L)) # initializing the permutation matrix if (cmp(init_type, 'random') == 0): print 'random initialization is being used ... \n' PI_0 = init_random(n_obs) elif (cmp(init_type, 'eig') == 0): print 'sorted eigenvector initialization is being used ... \n' PI_0 = init_eig(K, L, n_obs) else: print 'wrong initialization type ... ' # centering of kernel matrices H = numpy.eye(n_obs) - numpy.ones(n_obs) / n_obs K = numpy.dot(H, numpy.dot(K, H)) L = numpy.dot(H, numpy.dot(L, H)) print 'initial objective: ' print numpy.trace(numpy.dot(numpy.dot(numpy.dot(PI_0, K), PI_0.T), L)) # iterative linear assignment solution PI_t = numpy.zeros((n_obs, n_obs)) for i in xrange(n_iter): print 'iteration : ', i starttime = time.clock() grad = compute_gradient(K, L, PI_0) # L * P_0 * K stoptime = time.clock() print 'computing gradient takes %f seconds ' % (stoptime - starttime) # convert grad (profit matrix) to cost matrix # assuming it is a finite cost problem, thus # all the elements are substracted from the max value of the whole matrix cost_matrix = -grad starttime = time.clock() #indexes = LAPJV.lap(cost_matrix)[1] indexes = hungarian.hungarian(cost_matrix)[0] indexes = numpy.array(indexes) stoptime = time.clock() print 'lap solver takes %f seconds ' % (stoptime - starttime) PI_t = numpy.eye(n_obs) PI_t = PI_t[indexes, ] # convex combination PI = (1 - llambda) * PI_0 + (llambda) * PI_t # gradient ascent #PI = PI_0 + llambda*compute_gradient(K,L,PI_t) # computing the objective function obj_funct = numpy.trace( numpy.dot(numpy.dot(numpy.dot(PI_t, K), PI_t.T), L)) print 'objective function value : ', obj_funct print '\n' # another termination criteria if (numpy.trace(numpy.dot(numpy.dot(numpy.dot(PI, K), PI.T), L)) - numpy.trace(numpy.dot(numpy.dot(numpy.dot(PI_0, K), PI_0.T), L)) <= 1e-5): PI_final = PI_t break PI_0 = PI if (i == n_iter - 1): PI_final = PI_t return PI_final
def dr_cluster(data, method, gamma, params, clusters, stepsize, rows_toload, dropped_class_numbers): if (method == "Kmeans2D"): components = 2 if (method == "Kmeans1D" or method == "Thresholding"): components = 1 flag = 0 resetflag = 0 logger.writelog(components, "Components") logger.result_open(method) print(method) max_sc = -100.0 best_purity = 0.0 best_gamma = 0.0 serial_num = 0 try: for i in range(0, params + 1): transformer = KernelPCA(n_components=components, kernel='rbf', gamma=gamma) data_transformed = transformer.fit_transform(data) df = pd.DataFrame(data_transformed) df.to_csv(KPCA_output_path, index=False, header=None) del df gc.collect() if (method == "Thresholding"): if (flag == 0): os.system("cc c_thresholding_new.c") flag = 1 start = timeit.default_timer() os.system("./a.out " + str(clusters) + " " + str(rows_toload)) end = timeit.default_timer() thresholding_time = (end - start) sc = silhouette.silhouette(KPCA_output_path, Thresholding_paths[1]) groundtruth_distribution, temp_assignment_error_matrix, row_ind, col_ind, class_numbers, purity = hungarian.hungarian( 't', Thresholding_paths[0], clusters, rows_toload, dropped_class_numbers) logger.writeresult(i + 1, clusters, method, thresholding_time, gamma, sc, purity) #print(i+1,thresholding_time,gamma,sc,purity) if (i < params): if (sc > max_sc): max_sc = sc best_gamma = gamma best_purity = purity serial_num = i + 1 if (i == (params - 1)): gamma = best_gamma sc = max_sc purity = best_purity if (i == params): print(best_gamma, max_sc, best_purity) logger.writeresult(" ", " ", " ", " ", " ", " ", " ") logger.writeresult(serial_num, clusters, method, thresholding_time, best_gamma, max_sc, best_purity) logger.writeresult(" ", " ", " ", " ", " ", " ", " ") logger.writefinalresult(serial_num, clusters, method, thresholding_time, best_gamma, max_sc, best_purity) write_hungarian_result(best_gamma, clusters, groundtruth_distribution, temp_assignment_error_matrix, row_ind, col_ind, class_numbers, best_purity, method, params, stepsize, dropped_class_numbers) else: kmeans_time = kmeans.kmeans(KPCA_output_path, KMeans_paths[1], clusters) kmeans.groundtruth_distribution(KMeans_paths[1], KMeans_paths[0], datafiles_names[0], datafiles_names[2], clusters) sc = silhouette.silhouette(KPCA_output_path, KMeans_paths[1]) groundtruth_distribution, temp_assignment_error_matrix, row_ind, col_ind, class_numbers, purity = hungarian.hungarian( 'k', KMeans_paths[0], clusters, rows_toload, dropped_class_numbers) logger.writeresult(i + 1, clusters, method, kmeans_time, gamma, sc, purity) #print(i+1,kmeans_time,gamma,sc,purity) if (i < params): if (sc > max_sc): max_sc = sc best_gamma = gamma best_purity = purity serial_num = i + 1 if (i == (params - 1)): gamma = best_gamma sc = max_sc purity = best_purity if (i == params): print(best_gamma, max_sc, best_purity) logger.writeresult(" ", " ", " ", " ", " ", " ", " ") logger.writeresult(serial_num, clusters, method, kmeans_time, best_gamma, max_sc, best_purity) logger.writeresult(" ", " ", " ", " ", " ", " ", " ") logger.writefinalresult(serial_num, clusters, method, kmeans_time, best_gamma, max_sc, best_purity) write_hungarian_result(best_gamma, clusters, groundtruth_distribution, temp_assignment_error_matrix, row_ind, col_ind, class_numbers, best_purity, method, params, stepsize, dropped_class_numbers) if (i < (params - 1)): gamma = gamma + stepsize except (KeyboardInterrupt, SystemExit, Exception) as ex: ex_type, ex_value, ex_traceback = sys.exc_info() trace_back = traceback.extract_tb(ex_traceback) logger.writelog(str(ex_type.__name__), "Exception Type") logger.writelog(str(ex_value), "Exception Message") logger.writelog(str(trace_back), "Traceback") finally: logger.result_close()
hungarian_result = {} # hungarian_alt_result = {} for count in input_count: print("-" * 80) print("Input task count ", count) agent_task_input = generate_input(count) dmb_cost = dmb(agent_task_input) idmb_cost = idmb(agent_task_input) print("DMB", dmb_cost) print("IDMB", idmb_cost) dmb_result[count] = dmb_cost idmb_result[count] = idmb_cost cost_matrix_input = generate_cost_matrix(agent_task_input) hungarian_cost = hungarian(cost_matrix_input) print("Hungarian", hungarian_cost) hungarian_result[count] = hungarian_cost # generating plots colors = ['orange', 'red', 'blue'] plt.rcParams.update({'font.size': 14}) plt.plot(list(dmb_result.keys()), list(dmb_result.values()), colors[0], label='DMB') plt.plot(list(idmb_result.keys()), list(idmb_result.values()), colors[1], label='IDMB') plt.plot(list(hungarian_result.keys()), list(hungarian_result.values()), colors[2], label='Hungarian') plt.legend(loc='upper left', fancybox=False, shadow=False, ncol=4, prop={'size': 10}) plt.xlabel("Number of tasks/agents") plt.ylabel("Cost") plt.show()
print("Time: ", elapsed_time) print("\n__________________________\n") #Test-Input-3.txt start = time.clock() mst_prims('test-input-3.txt') elapsed_time = (time.clock() - start) print('Running Algo tps_mst_prims:') print("File: test-input-3.txt: ") print("Time: ", elapsed_time) print("\n__________________________\n") #Test-Input-4.txt start = time.clock() hungarian('test-input-4') elapsed_time = (time.clock() - start) print('Running Algo Hungarian:') print("File: test-input-4: ") print("Time: ", elapsed_time) print("\n__________________________\n") #Test-Input-5.txt start = time.clock() fast('test-input-5') elapsed_time = (time.clock() - start) print('Running Algo Fast:') print("File: test-input-5: ") print("Time: ", elapsed_time) print("\n__________________________\n")
(mu,S,priors,gamma,err) = kcluster0.gmm(x,k,weights=w,nreplicates=10) #print "mu = " + str(mu) #print "idx = " + str(idx) if __name__ == "__main__": for trial in range(10): print "trial = " + str(trial) (x,k,w) = initialize() state = random.get_state() if True: (mu0,S0,priors0,gamma0,err0) = kcluster0.gmm(x,k,weights=w,nreplicates=10) random.set_state(state) (mu,S,priors,gamma,err) = kcluster.gmm(x,k,weights=w,nreplicates=10) D = (num.tile(mu[:,0].reshape((k,1)),(1,k)) - mu0[:,0])**2 + \ (num.tile(mu[:,1].reshape((k,1)),(1,k)) - mu0[:,1])**2 (a1,a2) = hungarian.hungarian(D) if True or \ num.max(num.abs(mu - mu0[a1,:])) > .001 or \ num.max(num.abs(S-S0[:,:,a1])) > .001 or \ num.max(num.abs(priors - priors0[a1])) > .001 or \ num.max(num.abs(gamma-gamma0[:,a1])) > .001 or \ num.abs(err - err0) > .001: print "cluster approx cluster0["+str(a1)+"]" print "max(|mu - mu0|) = " + str(num.max(num.abs(mu - mu0[a1,:]))) print "max(|S - S0|) = " + str(num.max(num.abs(S-S0[:,:,a1]))) print "max(|priors - priors0|) = " + str(num.max(num.abs(priors - priors0[a1]))) print "max(|gamma - gamma0|) = " + str(num.max(num.abs(gamma-gamma0))) print "(err - err0)/err0 = " + str((err-err0)/err0) else:
def resetAssignments(self, state: 'State'): # Allocates one distinct box to each of the goal and # allocates these boxes to distinct agents. # Reset the variables containing the assignments state.goalBoxAssignments = defaultdict(list) state.agentBoxAssignments = defaultdict(list) # BOX - GOAL ASSIGNMENT SECTION # Create containers for box IDs and box coords, in a similar way as we # did with goals in the __init__ method boxIdsToBeCompleted = defaultdict(list) boxCoordsToBeCompleted = defaultdict(list) goalIdsToBeCompleted = defaultdict(list) goalCoordsToBeCompleted = defaultdict(list) goalDistancesToBeCompleted = defaultdict(list) for box in state.boxes: # Do not assign boxes already in goal if state.parent is not None: if box.id in state.parent.boxIdsCompleted: continue boxIdsToBeCompleted[box.letter.lower()].append(box.id) boxCoordsToBeCompleted[box.letter.lower()].append(box.coords) for idx, goal in enumerate(state.goals): # Do not assign boxes already in goal if state.parent is not None: if goal.id in state.parent.goalIdsCompleted: continue goalIdsToBeCompleted[goal.letter.lower()].append(goal.id) goalCoordsToBeCompleted[goal.letter.lower()].append(goal.coords) goalDistancesToBeCompleted[goal.letter.lower()].append(State.goalDistances[idx]) # DO IT FOR EACH GOAL/BOX LETTER for letter in goalIdsToBeCompleted: # Create shorthands for readable code (they're just references, # without extra memory consumption) boxCoordsLetter = boxCoordsToBeCompleted[letter] goalDistsLetter = goalDistancesToBeCompleted[letter] box_cnt = len(boxCoordsLetter) goal_cnt = len(goalDistsLetter) # Create distance matrix containing distances # between boxes and goals in the following way: # # box1 box2 box3 # goal1 1 5 3 # goal2 8 7 5 # goal3 2 10 7 distMatrix = [[goalDistsLetter[i][boxCoordsLetter[j][0]][boxCoordsLetter[j][1]] for j in range(box_cnt)] for i in range(goal_cnt)] # Assign a box to each of the goals with the Hungarian algorithm # # Get assignments in a # [(goalX, boxY), (goalY, boxZ), (goalZ, boxX)] format # where goalX..Z and box..Z refers to their corresponding row/col # numbers in the DISTANCE MATRIX assignments = hungarian(distMatrix) # Get the real IDs of the assigned goals and boxes idAssignments = [(goalIdsToBeCompleted[letter][assignment[0]], boxIdsToBeCompleted[letter][assignment[1]]) for assignment in assignments] # Write the result as the assignment of boxes and goals of letter "letter". state.goalBoxAssignments[letter] = idAssignments #print(state.goalBoxAssignments, file=sys.stderr, flush=True) # AGENT - BOX ASSIGNMENT SECTION # Create containers for box and agent data, in a similar fashion as in # the above step, this time grouped by COLOR. # For the AGENT-BOX assignments, use boxes only that have an associated # goal. boxIdsWithGoals = defaultdict(list) boxCoordsWithGoals = defaultdict(list) agentIds = defaultdict(list) agentCoords = defaultdict(list) # Work with boxes only if they have a corresponding goal # otherwise there's no point in pushing them with agents for letter in state.goalBoxAssignments: boxIdsWithGoals[letter] = [assignment[1] for assignment in state.goalBoxAssignments[letter]] # If a box is not assigned to a goal, discard it # Group box data by color for box in state.boxes: if box.id not in boxIdsWithGoals[box.letter.lower()]: continue boxIdsWithGoals[box.color].append(box.id) boxCoordsWithGoals[box.color].append(box.coords) # Group agent data by color for agent in state.agents: agentIds[agent.color].append(agent.number) agentCoords[agent.color].append(agent.coords) boxIdsWithAgents = [] # DO IT FOR EACH AGENT/BOX COLOR ... for color in agentCoords: # Create shorthands for readable code (they're just references, # without extra memory consumption) boxCoordsColor = boxCoordsWithGoals[color] agentCoordsColor = agentCoords[color] box_cnt = len(boxCoordsColor) # Create distance matrix for Hungarian algorithm # box1 box2 box3 # agent1 1 5 3 # agent2 8 7 5 # agent3 2 10 7 distMatrix = [state.get_distance_one_to_many(agentCoord, boxCoordsColor) for agentCoord in agentCoordsColor] # Get assignments in a # [(agentX, boxY), (agentY, boxZ), (agentZ, boxX)] format # where agentX..Z and box..Z refers to their corresponding row/col # numbers in the DISTANCE MATRIX assignments = hungarian(distMatrix) # Get the real IDs of the assigned agents and boxes idAssignments = [(agentIds[color][assignment[0]], boxIdsWithGoals[color][assignment[1]]) for assignment in assignments] # Collect the IDs of boxes that have an associated agent boxIdsWithAgents += [assignment[1] for assignment in idAssignments] # Write the result as the assignment of agents and boxes of this color state.agentBoxAssignments[color] = idAssignments # Remove the GOAL_BOX assignments that have no associated agent. for color in state.goalBoxAssignments: state.goalBoxAssignments[color] = [assignment for assignment in state.goalBoxAssignments[color] if assignment[1] in boxIdsWithAgents] '''