2 files changed, 24 insertions, 56 deletions

diff --git a/include/brain/neural_network.h b/include/brain/neural_network.h
index ccd65eb..22e5e42 100644
--- a/include/brain/neural_network.h
+++ b/include/brain/neural_network.h
@@ -3,6 +3,7 @@
 #include <memory>
 #include <istream>
 #include <vector>
+#include <map>
 #include <string>
 
 #include "brain/neural_node.h"
@@ -23,10 +24,10 @@ public:
     int numberOfOutputs() const { return _outputs.size(); }
 	
 private:
-    std::vector<std::shared_ptr<Sensor>> _sensors;
+    std::map<int, std::shared_ptr<Sensor>> _sensors;
     std::shared_ptr<BiasNode> _biasNode;
-    std::vector<std::shared_ptr<Neuron>> _hiddenNodes;
-    std::vector<std::shared_ptr<Neuron>> _outputs;
+    std::map<int, std::shared_ptr<Neuron>> _hiddenNodes;
+    std::map<int, std::shared_ptr<Neuron>> _outputs;
 
     void parseFile(std::istream &&file);
     
diff --git a/src/brain/neural_network.cpp b/src/brain/neural_network.cpp
index 30da46c..d6a5e15 100644
--- a/src/brain/neural_network.cpp
+++ b/src/brain/neural_network.cpp
@@ -7,11 +7,11 @@ NeuralNetwork::NeuralNetwork(std::istream &&networkConfigFile, int numberOfSenso
     
     for (int i=0; i<numberOfSensors; ++i)
     {
-	_sensors.push_back(std::make_shared<Sensor>(i));
+	_sensors[i] = std::make_shared<Sensor>(i);
     }
     for (int i=0; i<numberOfOutputs; ++i)
     {
-	_outputs.push_back(std::make_shared<Neuron>(i));
+	_outputs[i] = std::make_shared<Neuron>(i);
     }
 
     parseFile(std::move(networkConfigFile));
@@ -25,11 +25,11 @@ NeuralNetwork::NeuralNetwork(std::istream &&networkConfigFile, std::vector<bool>
     {
 	auto sensor = std::make_shared<Sensor>(i);
 	sensor->setActivation(sensorInitialValues.at(i) ? 1 : 0);
-	_sensors.push_back(sensor);
+	_sensors[i] = sensor;
     }
     for (int i=0; i<numberOfOutputs; ++i)
     {
-	_outputs.push_back(std::make_shared<Neuron>(i));
+	_outputs[i] = std::make_shared<Neuron>(i);
     }
 
     parseFile(std::move(networkConfigFile));
@@ -84,79 +84,46 @@ void NeuralNetwork::addLink(std::shared_ptr<NeuralNode> source, std::shared_ptr<
 
 std::shared_ptr<Sensor> NeuralNetwork::findOrAddSensor(int id)
 {
-    std::shared_ptr<Sensor> result;
-    for (auto node : _sensors)
+    bool sensorExists = _sensors.count(id) > 0;
+    if (!sensorExists)
     {
-	if (node->id() == id)
-	{
-	    result = node;
-	    break;
-	}
+	_sensors[id] = std::make_shared<Sensor>(id);
     }
-    if (!result)
-    {
-	result = std::make_shared<Sensor>(id);
-	_sensors.push_back(result);
-    }
-    return result;
+    return _sensors.at(id);
 }
 
 std::shared_ptr<Neuron> NeuralNetwork::findOrAddNeuron(int id)
 {
-    std::shared_ptr<Neuron> result;
-    for (auto node : _hiddenNodes)
-    {
-	if (node->id() == id)
-	{
-	    result = node;
-	    break;
-	}
-    }
-    if (result)
-    {
-	return result;
-    }
-    for (auto node : _outputs)
+    bool isOutput = _outputs.count(id) > 0;
+    if (isOutput)
     {
-	if (node->id() == id)
-	{
-	    result = node;
-	    break;
-	}
+	return _outputs.at(id);
     }
-    if (result)
+
+    bool hiddenNeuronExists = _hiddenNodes.count(id) > 0;
+    if (!hiddenNeuronExists)
     {
-	return result;
+	_hiddenNodes[id] = std::make_shared<Neuron>(id);
     }
-    
-    result = std::make_shared<Neuron>(id);
-    _hiddenNodes.push_back(result);
-    return result;
+    return _hiddenNodes.at(id);
 }
 
 void NeuralNetwork::setInput(int inputIndex, double activation)
 {
-    for (auto sensor : _sensors)
-    {
-	if (sensor->id() == inputIndex)
-	{
-	    sensor->setActivation(activation);
-	    break;
-	}
-    }
+    _sensors.at(inputIndex)->setActivation(activation);
 }
 
 int NeuralNetwork::findMaxOutputIndex() const
 {
     double currentMaxActivation = 0;
     int currentMaxIndex = 0;
-    for (auto output : _outputs)
+    for (std::pair<const int, std::shared_ptr<Neuron>> outputPair : _outputs)
     {
-	double activation = output->activation();
+	double activation = outputPair.second->activation();
 	if (activation >= currentMaxActivation)
 	{
 	    currentMaxActivation = activation;
-	    currentMaxIndex = output->id();
+	    currentMaxIndex = outputPair.first;
 	}
     }
     return currentMaxIndex;