1
#include <Rcpp.h>
2

3
using namespace Rcpp;
4

5

6
#include<vector>
7
#include<algorithm>
8

9
// [[Rcpp::export]]
10

11 4
  NumericVector rcpp_groupmean(IntegerVector cat_vec, NumericVector val_vec) {
12
    // init
13 4
    IntegerVector ids_vec=na_omit(sort_unique(cat_vec));
14 4
    IntegerVector levels_vec=match(cat_vec,ids_vec)-1;
15 4
    NumericVector sum_vec(ids_vec.size());
16 4
    NumericVector count_vec(ids_vec.size());
17 4
    std::vector<double> cpp_mean_vec;
18 4
    cpp_mean_vec.reserve(ids_vec.size());
19 4
    std::vector<int> cpp_ids_vec;
20 4
    cpp_ids_vec.reserve(ids_vec.size());
21 4
    std::vector<double> cpp_counts_vec;
22 4
    cpp_counts_vec.reserve(ids_vec.size());
23

24
    /// main
25
    // calculate sums
26 4
    for (int i=0; i<cat_vec.size(); ++i) {
27 4
      if (!IntegerVector::is_na(cat_vec[i]) && !NumericVector::is_na(val_vec[i])) {
28 4
        count_vec[levels_vec[i]]+=1.0;
29 4
        sum_vec[levels_vec[i]]+=val_vec[i];
30
      }
31
    }
32

33
    // calculate means
34 4
    for (int i=0; i<sum_vec.size(); ++i) {
35 4
      if (count_vec[i]>0) {
36 4
        cpp_mean_vec.push_back(sum_vec[i]/count_vec[i]);
37 4
        cpp_ids_vec.push_back(ids_vec[i]);
38 4
        cpp_counts_vec.push_back(count_vec[i]);
39
      }
40
    }
41 4
    cpp_mean_vec.shrink_to_fit();
42 4
    cpp_ids_vec.shrink_to_fit();
43 4
    cpp_counts_vec.shrink_to_fit();
44

45
    // exports
46 4
    NumericVector ret_vec = wrap(cpp_mean_vec);
47 4
    ret_vec.attr("ids") = wrap(cpp_ids_vec);
48 4
    ret_vec.attr("counts") = wrap(cpp_counts_vec);
49 4
    return(ret_vec);
50
  }

Read our documentation on viewing source code .

Loading