diff --git a/src/lib/cpp/cpu/connected_components.cc b/src/lib/cpp/cpu/connected_components.cc index 825818b..6ac5ad0 100644 --- a/src/lib/cpp/cpu/connected_components.cc +++ b/src/lib/cpp/cpu/connected_components.cc @@ -641,10 +641,23 @@ std::vector> connected_components(const std::string &base_p } void count_sizes(int64_t *__restrict__ img, std::vector &sizes, const int64_t n_labels, const int64_t size) { - #pragma omp parallel for schedule(static) - for (int64_t i = 0; i < size; i++) { - assert (img[i] <= n_labels && "Label out of bounds"); - sizes[img[i]]++; + std::vector> sizes_thread(omp_get_max_threads(), std::vector(n_labels+1, 0)); + #pragma omp parallel + { + int64_t thread_id = omp_get_thread_num(); + std::vector local_sizes = sizes_thread[thread_id]; + + #pragma omp for + for (int64_t i = 0; i < size; i++) { + local_sizes[img[i]]++; + } + + #pragma omp for + for (int64_t i = 0; i < n_labels+1; i++) { + for (int64_t j = 0; j < omp_get_num_threads(); j++) { + sizes[i] += sizes_thread[j][i]; + } + } } }