使用模板类实现KM算法

0.前言

      自己在写一个小程序时，遇到了一个类似于“完备匹配下的最大权匹配”的优化问题。在网上搜了下相关资料，了解到对应的匈牙利算法与KM算法，而且已经都有大神进行了详细讲解和代码的编写。唯一的不同之处是我参考的文章中KM算法目标是匹配结果最大为目标，而我的程序中是以匹配结果最小为目标。自己把代码改写了下，并封装为类。验证结果表明代码没有问题~

1.KM算法

      自己是面向工程应用的渣渣程序员，原理方面的参考网上大神讲解。有关KM算法原理和基本实现，可以参考如下三篇文章：

http://www.cnblogs.com/wenruo/p/5264235.html KM算法详解+模板

http://blog.csdn.net/kevinjqy/article/details/54584114 分配问题与匈牙利算法

http://blog.csdn.net/dark_scope/article/details/8880547 趣写算法系列之匈牙利算法

2.代码实现

      使用模板类实现了匹配结果最大为目标的KM算法。

KM_Template.h

#ifndef _KM_TEMPLATE_H_
#define _KM_TEMPLATE_H_

#include <vector>

template<class Type> class KM
{
public:
	KM();
	~KM();
	void init(int size, std::vector<std::vector<Type>> diff);
	bool dfs(int a);
	void KM_Calc(); 
	int* getMatch();
	Type getdiffSum();

private:
	std::vector<std::vector<Type>> diff;// 差异度矩阵(方阵)
	Type* ex_a;// 每个a类对象的期望
	Type* ex_b;// 每个b类对象的期望
	bool* vis_a;// 记录每一轮匹配过的a类对象
	bool* vis_b;// 记录每一轮匹配过的b类对象
	int* match;// 记录每一个b类对象匹配到的a类对象索引，如果没有匹配到则为-1
	Type* slack;// 记录每个b类对象能被a类对象匹配到所需要减少的最小差异度
	int size;// 记录差异度矩阵的边长
	Type diffSum;// 记录匹配结果的总体差异度大小
};

#endif

KM_Template.cpp

#include "KM_Template.h"

template<class Type> KM<Type>::KM()
{
	// 参数初始化
	ex_a = ex_b = NULL;
	slack = NULL;
	match = NULL;
	vis_a = vis_b = NULL;
	size = 0;
	diffSum = 0;
}

template<class Type> KM<Type>::~KM()
{
	/* 释放之前分配了的内存 */
	if(ex_a)
	{
		delete []ex_a;
		ex_a = NULL;
	}
	if(ex_b)
	{
		delete []ex_b;
		ex_b = NULL;
	}
	if(vis_a)
	{
		delete []vis_a;
		vis_a = NULL;
	}
	if(vis_b)
	{
		delete []vis_b;
		vis_b = NULL;
	}
	if(slack)
	{
		delete []slack;
		slack = NULL;
	}

	if(match)
	{
		delete []match;
		match = NULL;
	}
	
}

template<class Type> void KM<Type>::init(int size, std::vector<std::vector<Type>> diff)
{
	/* 获取差异度矩阵，根据矩阵的边长为其他变量分配空间 */
	this->diff = diff;
	this->size = size;
	ex_a = new Type[size];
	ex_b = new Type[size];
	vis_a = new bool[size];
	vis_b = new bool[size];
	slack = new Type[size];
	match = new int[size];
}

template<class Type> bool KM<Type>::dfs(int a)
{
	vis_a[a] = true;
	for(int b = 0; b < size; b++)
	{
		if(vis_b[b])// 每一轮匹配中b类每个对象只尝试一次
			continue;
		Type c = diff[a][b];
		Type gap = ex_a[a] + ex_b[b] - c;// 差异度符合要求

		if(fabs(float(gap)) < 1e-6)
		{
			vis_b[b] = true;
			if(match[b] == -1 || dfs(match[b]))// 匹配成功
			{
				match[b] = a;
				return true;
			}
		}
		else
		{
			slack[b] = std::max(slack[b], gap);// 匹配失败，记录该b类对象若想成功匹配，其差异度减小的最小值
		}
	}

	return false;
}

template<class Type> void KM<Type>::KM_Calc()
{
	for(int i = 0; i < size; i++)
	{
		ex_b[i] = 0;// 每个b类对象的差异度初始为0
		match[i] = -1;
		ex_a[i] = (diff)[i][0];
		for(int j = 1; j < size; j++)
		{
			ex_a[i] = std::min(ex_a[i], (diff)[i][j]);// 每个a类对象的初始期望值是与之相连的每个b类对象差异度的最小值
		}

	}

	for(int i = 0; i < size; i++)
	{
		for(int m = 0; m < size; m++)
			slack[m] = -1000000;
		while(1)
		{
			for(int m = 0; m <size; m++)
			{
				vis_a[m] = false;
				vis_b[m] = false;
			}
			if(dfs(i))
				break;// 匹配成功，退出
			// 匹配失败
			Type d = -1000000.0;
			for(int j = 0; j < size; j++)
			{
				if(!vis_b[j])
					d = std::max(d, slack[j]);// 注意这里d小于0
			}
			for(int j = 0; j < size; j++)
			{
				if(vis_a[j])
					ex_a[j] -= d;// 参加过匹配的a类对象只能提高差异度的期望值
				if(vis_b[j])
					ex_b[j] += d;// 参加过匹配的b类对象有资格“要求”a类对象降低差异度
				else
					slack[j] -= d;
			}
		}
	}
	
	diffSum = 0;
	// 计算匹配结果的总差异度
	for(int i = 0; i < size; i++)
	{
		diffSum += (diff)[match[i]][i];
	}
}

template<class Type> int* KM<Type>::getMatch()
{
	return match;
}

template<class Type> Type KM<Type>::getdiffSum()
{
	return diffSum;
}

      main.cpp提供调用示例，其中的数据来自http://blog.csdn.net/kevinjqy/article/details/54584114

#include <iostream>

#include "KM_Template.h"
#include "KM_Template.cpp"

int main()
{
	/* 构造差异度矩阵 */
	std::vector<std::vector<int>> diff;
	diff.resize(4);
	for(int i = 0; i < diff.size(); i++)
	{
		diff[i].resize(4);
	}
	diff[0][0] = 90;
	diff[0][1] = 75;
	diff[0][2] = 75;
	diff[0][3] = 80;
	diff[1][0] = 35;
	diff[1][1] = 85;
	diff[1][2] = 55;
	diff[1][3] = 65;
	diff[2][0] = 125;
	diff[2][1] = 95;
	diff[2][2] = 90;
	diff[2][3] = 105;
	diff[3][0] = 45;
	diff[3][1] = 110;
	diff[3][2] = 95;
	diff[3][3] = 115;

	/* 初始化算法对象 */
	KM<int> km;
	km.init(diff.size(), diff);
	/* 计算结果并输出 */
	km.KM_Calc();

	std::cout<<"diffSum: "<<km.getdiffSum()<<std::endl;
	int* match = km.getMatch();
	for(int i = 0; i < diff.size(); i++)
	{
		std::cout<<match[i]<<" ";
	}
	std::cout<<std::endl;

	return 0;
}