基于索引数组向量化拆分NumPy字符串数组为数值与单位数组

问题背景 #

给定原始字符串数组：

import numpy as np
og_arr = np.array([['5mm', '45"', '300 mm WT', 'Nan'], ['50mm', '3/5"', 'Nan', 'Nan']])

以及每个字符串中单位的起始索引数组：

index_arr = np.array([[1, 2, 4, -1], [2, 3, -1, -1]])

需要将og_arr拆分为数值数组values和单位数组units，预期结果如下：

values = [['5', '45', '300', 'Nan'], ['50', '3/5', 'Nan', 'Nan']]
units =  [['mm', '"', 'mm WT', ''], ['mm', '"', '', '']]

已实现循环解法，现寻求纯向量化方案。

解决方案 #

方法1：利用np.vectorize实现高效伪向量化

np.vectorize虽为伪向量化封装，但内部会优化执行逻辑，比手动循环高效，且代码简洁易维护：

import numpy as np

# 初始化输入数组
og_arr = np.array([['5mm', '45"', '300 mm WT', 'Nan'], ['50mm', '3/5"', 'Nan', 'Nan']])
index_arr = np.array([[1, 2, 4, -1], [2, 3, -1, -1]])

# 处理索引为-1的情况：替换为对应字符串长度，避免切片报错
str_lengths = np.vectorize(len)(og_arr)
adjusted_indices = np.where(index_arr == -1, str_lengths, index_arr)

# 定义拆分逻辑函数
def extract_value(s, idx):
    return 'Nan' if s == 'Nan' else s[:idx]

def extract_unit(s, idx):
    return '' if (s == 'Nan' and idx == str_lengths[np.where(og_arr == s)][0]) else s[idx:]

# 向量化执行拆分
values = np.vectorize(extract_value)(og_arr, adjusted_indices)
units = np.vectorize(extract_unit)(og_arr, adjusted_indices)

验证输出：

print(values)
# 输出：
# [['5' '45' '300' 'Nan']
#  ['50' '3/5' 'Nan' 'Nan']]

print(units)
# 输出：
# [['mm' '"' 'mm WT' '']
#  ['mm' '"' '' '']]

方法2：底层纯向量化实现（无np.vectorize）

如果需要完全不依赖循环封装，可利用NumPy的广播和字符串视图，但由于数组中字符串长度不一致，需先对齐长度，实现复杂度较高：

# 找到最长字符串长度
max_len = max(len(s) for row in og_arr for s in row)
# 将所有字符串填充到相同长度（用空字符补全）
padded_arr = np.char.ljust(og_arr, max_len)
# 将字符串转换为字符数组
char_arr = padded_arr.view('U1').reshape(og_arr.shape + (max_len,))

# 生成数值部分的掩码（索引小于adjusted_indices的位置）
value_mask = np.arange(max_len)[np.newaxis, np.newaxis, :] < adjusted_indices[..., np.newaxis]
# 生成数值数组：保留掩码为True的字符并拼接
values = np.array([''.join([c for c, mask in zip(row, row_mask) if mask]) 
                  for row, row_mask in zip(char_arr, value_mask)])
values = values.reshape(og_arr.shape)
# 还原Nan值
values[og_arr == 'Nan'] = 'Nan'

# 生成单位部分的掩码
unit_mask = np.arange(max_len)[np.newaxis, np.newaxis, :] >= adjusted_indices[..., np.newaxis]
# 生成单位数组：保留掩码为True的字符并拼接
units = np.array([''.join([c for c, mask in zip(row, row_mask) if mask]).strip() 
                  for row, row_mask in zip(char_arr, unit_mask)])
units = units.reshape(og_arr.shape)

该方法完全基于NumPy向量化操作，但代码繁琐，仅在对性能有极致要求时推荐使用。

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