基于first和zero索引的数组元素ID编号简洁实现方案

问题描述 #

需要为数组中的每个元素分配ID编号，示例如下：

# case 1 
#points              z            f 
#index:      0   1   2   3   4    5    6   7   8   9   10
#ID:        -2  -1   0   1   2   -8   -7  -6  -5  -4  -3

# case 2
#points              f            z 
#index:      0   1   2   3   4    5    6   7   8   9   10
#ID:         6   7  -3  -2  -1    0    1   2   3   4   5

其中zero（原变量z）是存储“零触发”索引的变量，first（原变量f）是存储“首个”元素索引的变量。要求用简洁的写法替代多分支if-else语句，适配以下代码框架：

n_samples = None
first  = None
zero = None

def index(i):
    ID = None
    
    # Your solution

    return ID
    
for i in range(0, n_samples):
    print(index(i))

简洁实现方案 #

通过观察ID生成规律，可用一行条件表达式完成计算，无需冗余分支：

n_samples = None
first  = None
zero = None

def index(i):
    return (i - zero) if (zero < first and i < first) or (first < zero and i >= first) else \
           (i - zero) - n_samples if zero < first else (i - zero) + n_samples
    
for i in range(0, n_samples):
    print(index(i))

规律说明 #

• 当zero < first（Case 1场景）：
  • 索引i小于first时，ID直接取i - zero，对应示例中0-4号索引的结果
  • 索引i大于等于first时，ID取(i - zero) - n_samples，对应示例中5-10号索引的结果
• 当first < zero（Case 2场景）：
  • 索引i大于等于first时，ID直接取i - zero，对应示例中2-10号索引的结果
  • 索引i小于first时，ID取(i - zero) + n_samples，对应示例中0-1号索引的结果

该写法完全适配两种场景，代码简洁且逻辑清晰。

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