diff --git a/main.cpp b/main.cpp
index 45300a0f6b82b4dfcfd0b6144f4deae96dfc8fdd..17af0a1f7d326006b813387c6d44e5bbdc0f7ae7 100644
--- a/main.cpp
+++ b/main.cpp
@@ -45,7 +45,33 @@ Matrix compute_dct_matrix(uint32_t n_samples) {
 }
 
 //--------------------------------------------------------------------------------------------------
-// Selects certain columns of a matrix.
+std::vector<uint32_t> select_subsample(uint32_t n_samples, uint32_t n_subsamples) {
+    std::vector<uint32_t> subset_indices;
+    subset_indices.resize(n_samples);
+    for (uint32_t i = 0; i < n_samples; ++i) {
+        subset_indices[i] = i;
+    }
+    auto rng = std::default_random_engine();
+    rng.seed(std::default_random_engine::result_type(84610373847));
+    std::shuffle(std::begin(subset_indices), std::end(subset_indices), rng);
+    subset_indices.resize(n_subsamples);
+    std::sort(subset_indices.begin(), subset_indices.end());
+    return subset_indices;
+}
+
+//--------------------------------------------------------------------------------------------------
+// Selects certain elements of a Vector.
+Vector vector_subset(Vector const& vector, std::vector<uint32_t> const& subset_indices) {
+    Vector subset(subset_indices.size());
+    for (uint32_t i = 0; i < subset_indices.size(); ++i) {
+        auto const index = subset_indices[i];
+        subset[i] = vector[index];
+    }
+    return subset;
+}
+
+//--------------------------------------------------------------------------------------------------
+// Selects certain columns of a Matrix.
 Matrix matrix_subset(Matrix const& matrix, std::vector<uint32_t> const& subset_indices) {
     Matrix subset(matrix.rows(), subset_indices.size());
     for (uint32_t i = 0; i < subset_indices.size(); ++i) {
@@ -79,24 +105,9 @@ int main() {
 
     // Part (d)
     constexpr uint32_t n_subsamples = 100;
-    std::vector<uint32_t> subset_indices;
-    subset_indices.resize(n_samples);
-    for (uint32_t i = 0; i < n_samples; ++i) {
-        subset_indices[i] = i;
-    }
-    auto rng = std::default_random_engine();
-    rng.seed(std::default_random_engine::result_type(84610373847));
-    std::shuffle(std::begin(subset_indices), std::end(subset_indices), rng);
-    subset_indices.resize(n_subsamples);
-    std::sort(subset_indices.begin(), subset_indices.end());
-    Vector subset_sample_times(n_subsamples);
-    Vector subset_sample_values(n_subsamples);
-    for (uint32_t i = 0; i < n_subsamples; ++i) {
-        auto const index = subset_indices[i];
-        subset_sample_times[i] = sample_times[index];
-        subset_sample_values[i] = sample_values[index];
-        //std::cout << index << '\n';
-    }
+    std::vector<uint32_t> const subset_indices = select_subsample(n_samples, n_subsamples);
+    Vector subset_sample_times = vector_subset(sample_times, subset_indices);
+    Vector subset_sample_values = vector_subset(sample_values, subset_indices);
     python_print("subset_sample_times", subset_sample_times);
     python_print("subset_sample_values", subset_sample_values);
 
@@ -104,11 +115,7 @@ int main() {
     Matrix const subset_dct_matrix = matrix_subset(dct_matrix, subset_indices);
     Vector recovered_dct = Vector::Random(n_samples);
     recovered_sample_values = dct_matrix.transpose() * recovered_dct;
-    Vector subset_recovered_sample_values(n_subsamples);
-    for (uint32_t i = 0; i < n_subsamples; ++i) {
-        auto const index = subset_indices[i];
-        subset_recovered_sample_values[i] = recovered_sample_values[index];
-    }
+    Vector subset_recovered_sample_values = vector_subset(recovered_sample_values, subset_indices);
     Vector subset_differences = subset_sample_values - subset_recovered_sample_values;
     Scalar loss = subset_differences.squaredNorm();
     Vector gradient = -2 * subset_dct_matrix * subset_differences;