Skip to content

一道与 factor 有关的问题

W 问了个问题,

问题:对于一个4维向量A=(a1,a2,a3,a4),我想要将其转化为另一个4维向量B=(b1,b2,b3,b4),b1=1,然后对于b_i,如果a_i在A的前i-1个元素里没出现过,那么b_i=max(b[1:(i-1)])+1,否则b_i=出现过的那个元素a_k对应的b_k。 例子:(23,12,23,13) → (1,2,1,3), (20,20,20,20)→(1,1,1,1)

因为逻辑挺清晰的,所以第一感觉便是自己写函数,于是随手写了个 Julia 版本的,

function f(A)
    visited = Dict{Any, Int}()
    n = length(A)
    B = zeros(Int, n)
    count = 0
    for i = 1:n
        if get(visited, A[i], 0) > 0 # if A[i] has been visited, return the index, otherwise return 0
            B[i] = visited[A[i]]
        else
            count += 1
            B[i] = count
            visited[A[i]] = B[i]
        end
    end
    return B
end

其中 Dict 主要是为了缩小搜索空间。不过 W 指出在 R 中自己写函数挺费时的,然后他自己想出了个很优雅的解法

ff <- function(A) { as.numeric(factor(A,levels = unique(A))) }

因为自己想复习下 Rcpp,所以将上述 Julia 代码改写成 Rcpp,

#include <Rcpp.h>
#include <map>
using namespace Rcpp;

// [[Rcpp::export]]
IntegerVector f(NumericVector A) {
    std::map<float, int> visited;
    int count = 0;
    int n = A.size();
    IntegerVector B(n);
    for(int i = 0; i < n; i++) {
        if (visited.find(A[i]) == visited.end()) {
            count += 1;
            B[i] = count;
            visited[A[i]] = B[i];
        } else {
            B[i] = visited[A[i]];
        }
    }
    return B;
}

/*** R
f(c(20, 20, 20, 20))
*/

然后与 W 的方法进行比较,部分测试结果如下,

可见 Rcpp 版本 f 略优于 W 的方法 ff。另外,也将其与 Julia 版本进行比较,

image

有点惊讶,速度竟然快这么多!