1
#include <Rcpp.h>
2

3
using namespace Rcpp;
4

5

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

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

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

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

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

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

Read our documentation on viewing source code .

Loading