在R中优化马尔可夫链全状态序列及概率生成代码性能

原始场景 #

马尔可夫链构建代码 #

set.seed(123)

n_states <- 5

matrix <- matrix(runif(n_states^2), nrow=n_states)

# 设置部分转移概率为0
matrix[1, 4:5] <- 0 
matrix[5, 1:3] <- 0  
matrix[2, 5] <- 0    

transition_matrix <- t(apply(matrix, 1, function(x) x/sum(x)))

rownames(transition_matrix) <- paste0("S", 1:n_states)
colnames(transition_matrix) <- paste0("S", 1:n_states)

print(round(transition_matrix, 3))

生成的转移矩阵 #

S1         S2        S3         S4        S5
S1 0.2229340 0.03531601 0.7417500 0.00000000 0.0000000
S2 0.3910569 0.26197885 0.2248868 0.12207747 0.0000000
S3 0.1536622 0.33530265 0.2545791 0.01580275 0.2406533
S4 0.2652280 0.16563210 0.1719994 0.09849610 0.2986444
S5 0.0000000 0.00000000 0.0000000 0.59278229 0.4072177

原始递归枚举函数 #

# 生成不同步数的状态序列函数
find_sequences_all_turns <- function(transition_matrix, start_state = 1, max_turns = 5) {
    n_states <- nrow(transition_matrix)
    
    all_sequences <- list()
    all_probabilities <- numeric()
    all_turns <- numeric()
    seq_counter <- 1
    
    generate_sequence <- function(current_seq, current_prob, steps_left, total_steps) {
        if(length(current_seq) > 1) {
            all_sequences[[seq_counter]] <<- current_seq
            all_probabilities[seq_counter] <<- current_prob
            all_turns[seq_counter] <<- total_steps - steps_left
            seq_counter <<- seq_counter + 1
        }
        
        if(steps_left == 0) {
            return()
        }
        
        current_state <- current_seq[length(current_seq)]
        possible_next_states <- which(transition_matrix[current_state,] > 0)
        
        for(next_state in possible_next_states) {
            prob <- transition_matrix[current_state, next_state]
            generate_sequence(
                c(current_seq, next_state),
                current_prob * prob,
                steps_left - 1,
                total_steps
            )
        }
    }
    
    generate_sequence(c(start_state), 1, max_turns - 1, max_turns)
    
    result_df <- data.frame(
        turn = all_turns,
        sequence_no = 1:length(all_sequences),
        sequence = sapply(all_sequences, paste, collapse=""),
        probability = all_probabilities
    )
    
    result_df <- result_df[order(result_df$turn, -result_df$probability),]
    rownames(result_df) <- NULL
    
    return(result_df)
}

调用结果示例 #

sequences_df <- find_sequences_all_turns(transition_matrix)

> sequences_df
    turn sequence_no sequence  probability
1      2         154       13 7.417500e-01
2      2           1       11 2.229340e-01
3      2          65       12 3.531601e-02
4      3         171      132 2.487108e-01
5      3         193      133 1.888341e-01
6      3         243      135 1.785046e-01
7      3          40      113 1.653613e-01
8      3         155      131 1.139789e-01
9      3           2      111 4.969955e-02
10     3          66      121 1.381057e-02
11     3         218      134 1.172169e-02
12     3          82      122 9.252048e-03
13     3         104      123 7.942105e-03
14     3          18      112 7.873137e-03
15     3         129      124 4.311289e-03

概率验证代码 #

library(dplyr)

probability_sums <- sequences_df %>%
    group_by(turn) %>%
    summarise(
        total_probability = sum(probability),
        num_sequences = n(),
        check_sum_to_one = abs(total_probability - 1) < 1e-10
    )

print(probability_sums)

性能优化方案 #

原始递归函数的瓶颈在于：递归调用的栈开销、全局变量的频繁修改（<<-）、动态列表的内存扩容、以及序列拼接时的对象复制。以下是针对性的优化方法：

1. 迭代法生成序列，预分配内存 #

用迭代替代递归，提前计算每一步的序列数量，预分配内存避免动态扩容的开销。

find_sequences_iterative <- function(transition_matrix, start_state = 1, max_turns = 5) {
    n_states <- nrow(transition_matrix)
    # 预计算每个状态的可达下一个状态及对应概率
    next_states_list <- lapply(1:n_states, function(s) {
        idx <- which(transition_matrix[s, ] > 0)
        list(states = idx, probs = transition_matrix[s, idx])
    })
    
    # 初始化：步数为1的序列（起始状态）
    current_sequences <- list(c(start_state))
    current_probs <- c(1)
    result_list <- list()
    
    for(turn in 2:max_turns) {
        new_sequences <- list()
        new_probs <- numeric()
        seq_idx <- 1
        
        # 遍历当前所有序列，生成下一步序列
        for(i in seq_along(current_sequences)) {
            last_state <- current_sequences[[i]][length(current_sequences[[i]])]
            next_info <- next_states_list[[last_state]]
            
            # 批量生成新序列
            for(j in seq_along(next_info$states)) {
                new_sequences[[seq_idx]] <- c(current_sequences[[i]], next_info$states[j])
                new_probs[seq_idx] <- current_probs[i] * next_info$probs[j]
                seq_idx <- seq_idx + 1
            }
        }
        
        # 存入结果
        result_list[[turn]] <- data.frame(
            turn = turn,
            sequence = sapply(new_sequences, paste, collapse = ""),
            probability = new_probs
        )
        
        # 更新当前序列用于下一步迭代
        current_sequences <- new_sequences
        current_probs <- new_probs
    }
    
    # 合并所有结果并排序
    result_df <- do.call(rbind, result_list)
    result_df$sequence_no <- 1:nrow(result_df)
    result_df <- result_df[order(result_df$turn, -result_df$probability), ]
    rownames(result_df) <- NULL
    # 调整列顺序
    result_df <- result_df[, c("turn", "sequence_no", "sequence", "probability")]
    
    return(result_df)
}

2. 向量化+矩阵运算优化概率计算 #

利用矩阵幂运算快速计算各序列的概率，同时用笛卡尔积生成所有可能序列（注意过滤不可达的序列）。

find_sequences_vectorized <- function(transition_matrix, start_state = 1, max_turns = 5) {
    n_states <- nrow(transition_matrix)
    state_labels <- 1:n_states
    
    # 生成所有步数的序列模板，过滤不可达的组合
    generate_valid_sequences <- function(k) {
        # 生成k步的所有可能状态组合
        all_combs <- expand.grid(rep(list(state_labels), k))
        # 过滤掉转移概率为0的序列
        valid <- apply(all_combs, 1, function(seq) {
            all(sapply(1:(k-1), function(i) {
                transition_matrix[seq[i], seq[i+1]] > 0
            }))
        })
        valid_combs <- all_combs[valid, ]
        # 转为字符串序列
        sequences <- apply(valid_combs, 1, paste, collapse = "")
        # 计算每个序列的概率
        probs <- apply(valid_combs, 1, function(seq) {
            prod(sapply(1:(k-1), function(i) {
                transition_matrix[seq[i], seq[i+1]]
            }))
        })
        # 只保留以start_state开头的序列
        start_mask <- valid_combs[, 1] == start_state
        data.frame(
            turn = k,
            sequence = sequences[start_mask],
            probability = probs[start_mask]
        )
    }
    
    # 生成2到max_turns的所有有效序列
    result_list <- lapply(2:max_turns, generate_valid_sequences)
    result_df <- do.call(rbind, result_list)
    result_df$sequence_no <- 1:nrow(result_df)
    result_df <- result_df[order(result_df$turn, -result_df$probability), ]
    rownames(result_df) <- NULL
    result_df <- result_df[, c("turn", "sequence_no", "sequence", "probability")]
    
    return(result_df)
}

3. 使用data.table提升数据处理效率 #

data.table在批量数据操作和排序上比data.frame快很多，适合处理大量序列数据：

library(data.table)

find_sequences_datatable <- function(transition_matrix, start_state = 1, max_turns = 5) {
    n_states <- nrow(transition_matrix)
    next_states_list <- lapply(1:n_states, function(s) {
        idx <- which(transition_matrix[s, ] > 0)
        data.table(next_state = idx, prob = transition_matrix[s, idx])
    })
    
    # 初始化：步数1的序列
    dt <- data.table(sequence = as.character(start_state), probability = 1, turn = 1)
    
    result_dt <- data.table()
    
    for(turn in 2:max_turns) {
        # 连接下一个状态
        current_dt <- dt[turn == turn - 1]
        # 扩展每个序列的下一个状态
        expanded <- current_dt[, next_states_list[[as.integer(substr(sequence, nchar(sequence), nchar(sequence))))]], by = .(sequence, probability)]
        # 生成新序列和概率
        expanded[, `:=`(
            sequence = paste0(sequence, next_state),
            probability = probability * prob,
            turn = turn
        )]
        expanded[, c("next_state", "prob") := NULL]
        
        # 加入结果
        result_dt <- rbind(result_dt, expanded)
        # 更新dt用于下一步
        dt <- rbind(dt, expanded)
    }
    
    # 添加sequence_no并排序
    result_dt[, sequence_no := .I]
    result_dt <- result_dt[order(turn, -probability)]
    setcolorder(result_dt, c("turn", "sequence_no", "sequence", "probability"))
    
    return(result_dt)
}

4. 避免字符串拼接，用整数向量存储序列 #

如果不需要字符串格式的序列，可以直接用整数向量存储，减少字符串操作的开销：

find_sequences_integer <- function(transition_matrix, start_state = 1, max_turns = 5) {
    n_states <- nrow(transition_matrix)
    next_states_list <- lapply(1:n_states, function(s) {
        list(states = which(transition_matrix[s, ] > 0), probs = transition_matrix[s, which(transition_matrix[s, ] > 0)])
    })
    
    current_seqs <- list(c(start_state))
    current_probs <- c(1)
    result_list <- list()
    
    for(turn in 2:max_turns) {
        new_seqs <- list()
        new_probs <- numeric()
        idx <- 1
        
        for(i in seq_along(current_seqs)) {
            last_state <- current_seqs[[i]][length(current_seqs[[i]])]
            next_info <- next_states_list[[last_state]]
            
            for(j in seq_along(next_info$states)) {
                new_seqs[[idx]] <- c(current_seqs[[i]], next_info$states[j])
                new_probs[idx] <- current_probs[i] * next_info$probs[j]
                idx <- idx + 1
            }
        }
        
        result_list[[turn]] <- data.frame(
            turn = turn,
            sequence = I(new_seqs),  # 用I()存储整数向量
            probability = new_probs
        )
        
        current_seqs <- new_seqs
        current_probs <- new_probs
    }
    
    result_df <- do.call(rbind, result_list)
    result_df$sequence_no <- 1:nrow(result_df)
    result_df <- result_df[order(result_df$turn, -result_df$probability), ]
    rownames(result_df) <- NULL
    result_df <- result_df[, c("turn", "sequence_no", "sequence", "probability")]
    
    return(result_df)
}

优化效果说明 #

• 迭代法避免了递归调用的栈开销，预分配内存减少了动态扩容的时间
• 向量化操作利用R的底层C实现，比显式循环更快
• data.table的连接和排序操作效率远高于data.frame
• 整数向量存储序列省去了字符串拼接的开销，适合后续需要进一步处理序列的场景

内容的提问来源于stack exchange，提问作者farrow90