如何在JIT中获取值的索引位置？及指定JIT函数代码的替代方案

1. 如何使用JIT获取某个值的索引位置？ #

在JIT编译环境（比如Numba）中，Python原生的list.index()方法通常不被支持，核心解决思路是手动实现遍历查找逻辑：

• 遍历目标序列的每个元素，逐一与目标值对比
• 找到匹配项时返回对应索引
• 未找到则返回特定标记（如-1）

示例代码（基于Numba JIT）：

from numba import jit

@jit(nopython=True)
def find_index(arr, target):
    for idx, val in enumerate(arr):
        # 针对序列类型元素，需逐个元素对比
        if all(v == t for v, t in zip(val, target)):
            return idx
    return -1

2. j = indices.index(list(neighbor))的可行替代方案 #

原代码依赖Python列表的index()方法，且indices为列表的列表（不可哈希，无法用partial），JIT环境下可采用以下两种方案：

方案一：手动遍历匹配 #

直接在JIT函数内实现遍历对比逻辑，适配你的场景：

from numba import jit

@jit(nopython=True)
def get_neighbor_index(indices, neighbor):
    for j, item in enumerate(indices):
        match = True
        for i in range(len(item)):
            if item[i] != neighbor[i]:
                match = False
                break
        if match:
            return j
    return -1

# 调用示例
indices = [[1,2], [4,5], [2,3]]
neighbor = (2,3)
j = get_neighbor_index(indices, neighbor)  # 返回2

方案二：提前构建哈希映射（预处理） #

若indices固定不变，可在JIT函数外提前将内部列表转为可哈希的元组，构建字典映射提升查找效率：

# 预处理：将列表转元组作为键，索引作为值
index_map = {tuple(item): idx for idx, item in enumerate(indices)}

from numba import jit

@jit(nopython=True)
def get_neighbor_index_via_map(index_map, neighbor):
    return index_map.get(neighbor, -1)

# 调用示例
j = get_neighbor_index_via_map(index_map, neighbor)  # 返回2

注意：若indices在JIT执行过程中动态变化，优先选方案一；若indices固定，方案二的查找效率更高。

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