Solver for Energy Balance

qq.cpp

// GistID: c61ccc4fe17d876defa7
#include <iostream>
#include <fstream>
#include <vector>
#include <math.h>
#include <algorithm>

using namespace std;

template <typename T>
void coutVector(vector<T> a){
	for(unsigned short i=0;i<a.size();i++){
		cout << a[i] << " ";
	}
	cout << endl;
	return;
}

template <typename T>
void coutVector2d(vector<vector<T> > a){
	for(unsigned short j=0;j<a.size();j++){
		for(unsigned short i=0;i<a[j].size();i++){
			cout << a[j][i] << " ";
		}
		cout << endl;
	}
	return;
}

bool checkEnd(vector<short> &currentBalance,unsigned short a){
	for(unsigned short i=0;i<currentBalance.size();i++)
		if(currentBalance[i] != a+i-currentBalance.size()+1)
			return 0;
	return 1;
}

bool checkEndAddCombination(vector<short> &currentBalance,unsigned short a){
	for(unsigned short i=0;i<currentBalance.size();i++)
		if(currentBalance[i] != a-i)
			return 1;
	return 0;
}

bool balanceRepeat(vector<short> balanceCombinations,short a){
	vector<bool> check(a,0);
	for(unsigned short j=0;j<balanceCombinations.size();j++){
		if(check[balanceCombinations[j]])
			return 1;
		check[balanceCombinations[j]] = 1;
	}
	return 0;
}

bool balanceCheck(vector<short> energy,vector<short> currentBalance, short neededBalance){
	for(unsigned short i=0;i<currentBalance.size();i++)
		neededBalance -= energy[currentBalance[i]];
	return neededBalance;
}

void addCombination(vector<short> currentBalance,vector<vector<short> > &balanceCombinationsNeeded){
	do
		balanceCombinationsNeeded.push_back(currentBalance);
	while(next_permutation(currentBalance.begin(),currentBalance.end()));
	return;
}

bool balanceChecker(vector<vector<vector<short> > > &balanceCombinations,vector<vector<short> > &balance,unsigned short a,unsigned short aCounter,vector<short> &mask){
	vector<short> prev = mask;
	for(unsigned short i=0;i<balanceCombinations[a][aCounter].size();i++){
		if(mask[balance[a][i+1]] == -1)
			mask[balance[a][i+1]] = balanceCombinations[a][aCounter][i];
		else if(mask[balance[a][i+1]] != balanceCombinations[a][aCounter][i]){
			mask = prev;
			return 1;
		}
	}
	vector<bool> lastCheck(mask.size(),0);
	for(unsigned short i=0;i<mask.size();i++){
		if(mask[i] != -1){
			if(lastCheck[mask[i]]){
				mask = prev;
				return 1;
			}
			lastCheck[mask[i]] = 1;
		}
	}
	return 0;
}

bool solve(vector<vector<vector<short> > > &balanceCombinations,vector<unsigned long long> &counter,vector<vector<short> > &balance, unsigned short step,unsigned short n,vector<short> &mask){
	vector<short> prev = mask;
	if(step == balance.size())
		return 1;
	for(unsigned long long i=0;i<balanceCombinations[step].size();i++){
		mask = prev;
		counter[step] = i;
		if(balanceChecker(balanceCombinations,balance,step,counter[step],mask))
			continue;
		if(solve(balanceCombinations,counter,balance,step+1,n,mask))
			return 1;
	}
	return 0;
}

int main()
{
	//    ifstream cin("input.txt");
	cout << "Number of energies: ";
	unsigned short n,m;
	cin >> n;
	vector<short> energy(n,0);
	for(unsigned short i=0;i<n;i++){
		cout << i << " energy: ";
		cin >> energy[i];
	}
	coutVector(energy);
	cout << "Number of balances: ";
	cin >> m;
	vector<vector<short> > balance(m,vector<short>(0,0));
	unsigned short balanceSize,balanceN;
	for(unsigned short i=0;i<m;i++){
		cout << i+1 << " balance: " << endl;
		cout << "Energy needed in balance: ";
		cin >> balanceSize;
		balance[i].push_back(balanceSize);
		cout << "Size: ";
		cin >> balanceSize;
		for(unsigned short j=1;j<=balanceSize;j++){
			cout << j << " number of energy cell: ";
			cin >> balanceN;
			balance[i].push_back(balanceN);
		}
	}
	cout << endl;
	vector<vector<vector<short> > > balanceCombinations(m,vector<vector<short> > (0,vector<short> (0,0)));
	vector<short> currentBalance(0,0);
	for(unsigned short i=0;i<m;i++){
		currentBalance.clear();
		currentBalance.resize(balance[i].size()-1);
		for(unsigned short j=1;j<currentBalance.size();j++)
			currentBalance[j] = j;
		currentBalance[currentBalance.size()-1]--;
		while(!checkEnd(currentBalance,energy.size()-1)){
			currentBalance[currentBalance.size()-1]++;
			for(short j=currentBalance.size()-1;j>0;j--){
				currentBalance[j-1] += currentBalance[j]/energy.size();
				currentBalance[j] %= energy.size();
			}
			for(short j=0;j<currentBalance.size();j++)
				if(currentBalance[j] >= currentBalance[j+1])
					currentBalance[j+1] = currentBalance[j]+1;
			if(balanceCheck(energy,currentBalance,balance[i][0]))
				continue;
			addCombination(currentBalance,balanceCombinations[i]);
		}
	}

	for(unsigned short i=0;i<m;i++){
		unsigned short imin = i;
		for(unsigned short j=i+1;j<m;j++){
			if(balanceCombinations[j].size() < balanceCombinations[imin].size())
				imin = j;
		}
		swap(balance[i],balance[imin]);
		swap(balanceCombinations[i],balanceCombinations[imin]);
	}

	for(unsigned short i=0;i<m;i++){
		cout << balanceCombinations[i].size() << " ";
	}
	cout << endl;

	vector<unsigned long long> counter(m,0);
	vector<short> mask(n,-1);
	solve(balanceCombinations,counter,balance,0,n,mask);
	vector<bool> used(n,0);
	for(unsigned short i=0;i<n;i++){
		used[mask[i]] = 1;
	}
	m = 0;
	for(unsigned short i=0;i<n;i++){
		if(mask[i] == -1){
			while(used[m])
				m++;
			mask[i] = m;
			m++;
		}
		cout << i << " - " << energy[mask[i]] << endl;
	}
	while(true){
		cout << "END OF THE PROGRAM. CLOSE IT WITH CLICKING ON RED CROSS IN TOP RIGHT CORNER!" << endl;
		getchar();
	}
	return 0;
}