diff --git a/lib/nBodySim.hpp b/lib/nBodySim.hpp
index a2f12931d08795f9677bdc550572ebda6c6e553a..d3db7a2e7d56967308a976a404f98de27354c4db 100644
--- a/lib/nBodySim.hpp
+++ b/lib/nBodySim.hpp
@@ -6,40 +6,40 @@
 
 class nBodySim {
 public:
-    // constructor reads in data from datafile
+    // Constructor reads in data from datafile
     nBodySim(std::string datafile);
-    // destructor deletes arrays 
+    // Destructor deletes arrays 
     ~nBodySim();
 
-    // run simulation for nSteps time steps of size dt, integration with leapfrog algorithm
+    // Run simulation for nSteps time steps of size dt, integration with leapfrog algorithm
     void runSimulation(double dt, unsigned nSteps);
 
-    // loops over all pairs of particles and calculates forces, calculates current mean interparticle distance 
+    // Loops over all pairs of particles and calculates forces, calculates current mean interparticle distance 
     void calculateForces();
-    // use treecode and multipole expansion to calculate forces
+    // Use treecode and multipole expansion to calculate forces
     void treeCalculateForces();
 
-    // loop over all pairs of particles and calculate mean interparticle distance
+    // Loop over all pairs of particles and calculate mean interparticle distance
     double calculateMeanInterparticleDistance();
-    // loop over all force vectors and calculate mean force magnitude
+    // Loop over all force vectors and calculate mean force magnitude
     double calculateMeanForceMagnitude();
 
-    // integrate using leapfrog algorithm. doesn't update forces, doesn't update tree
+    // Integrate using leapfrog algorithm. doesn't update forces, doesn't update tree
     void doTimeStep(double dt);
-    // update tree
+    // Update tree
     void updateTree();
 
-    // writes data to file, every row is a particle
+    // Writes data to file, every row is a particle
     void write2file(const double* array, std::string filename, unsigned dim) const;
-    // write current state of simulation to file
+    // Write current state of simulation to file
     void saveState2file(unsigned step, std::string filename) const;
 
-    // functions to get private variables
+    // Functions to get private variables
     double* getMasses() { return masses_; }
     double* getPositions() { return positions_; }
     double* getForces() { return forces_; }
     unsigned getNParticles() const { return nParticles_; }
-    
+
     double getMeanInterparticleDistance() const { return meanInterparticleDistance_; }
     double getMeanForceMagnitude() const { return meanForceMagnitude_; }
 
diff --git a/lib/node.hpp b/lib/node.hpp
index f470642e970e55d609e01bf0452bc3cece02fbe5..d8c900d0987504a50981f505b1762ca7cbad9547 100644
--- a/lib/node.hpp
+++ b/lib/node.hpp
@@ -4,25 +4,28 @@
 #include <vector>
 
 class Node {
+    using List = std::vector<unsigned>;
 public:
-    // default constructor creates a leaf node
+    // Default constructor creates a leaf node
     Node() = default;
-    // whoever creates the node is responsible for deciding which particles are in which octant and for allocating the localParticles array 
-    // node will write to localParticles array
+    // Whoever creates the node is responsible for deciding which particles are in which octant and for allocating the localParticles array 
+    // Node will write to localParticles array
     Node(Node* root, Node* parent, double* masses, double* particles, unsigned nParticles, unsigned depth, 
-        double size, double center[3], std::vector<unsigned>& localParticles, unsigned nLocalParticles);
-    // destuctor
+        double size, double center[3], List& localParticles, unsigned nLocalParticles);
+    // Destuctor
     ~Node();
-    // copy constructor (default)
+    // Copy constructor (default)
     Node(const Node&) = default;
-    // copy assignment
+    // Copy assignment
     Node& operator=(const Node&);
-    // update tree: visit each node, check if it needs to be split or merged (for example if a particle has left the region)
-    void update(std::vector<unsigned>& allParticles);
-    // calculate force on a particle recursively
+    // Update tree: visit each node, check if it needs to be split or merged (for example if a particle has left the region)
+    void update(List& allParticles);
+    // Calculate force on a particle recursively
     double* calculateForce(unsigned particle);
-    // calculate center of mass
+    // Calculate center of mass
     void calculateCenterOfMass();
+    // Function that takes a particle and returns the octant it is in. This function doesn't check if the particle is actually in the node.
+    unsigned getOctant(unsigned particle);
 
 private:
     Node* root_;
@@ -38,7 +41,8 @@ private:
     double center_[3];
     double centerOfMass_[3];
 
-    std::vector<unsigned> localParticles_;
+    // Contains the indices of the particles that are in this node for the particles_ array
+    List localParticles_;
     unsigned nLocalParticles_;
 
     double mass_;
diff --git a/lib/tree.hpp b/lib/tree.hpp
index cc4474dfd375475d731a6cb82034a010f8152889..bdab71c81f3ceebd39860867cf83b370dd0ad7dd 100644
--- a/lib/tree.hpp
+++ b/lib/tree.hpp
@@ -7,14 +7,14 @@
 
 class Tree {
 public:
-    // constructor
+    // Constructor
     Tree() = delete;
     Tree(double* masses, double* particles, double* forces_, unsigned nParticles, double size, double* center);
-    // destructor
+    // Destructor
     ~Tree();
-    // calculate force on a particle
+    // Calculate force on a particle
     void calculateForce(unsigned particle);
-    // update tree
+    // Update tree
     void update();
 
 private:
@@ -27,7 +27,7 @@ private:
 
     double size_;
     double center_[3];
-    
+
     std::vector<unsigned> allParticles_;
 };