python - 用 numpy 函数（如 numpy.where）替换多个 if elif

Question

Radio_index、n_x 和 n_y 是整数我写了一个可以运行的 if/elif 代码。目的是找到 radio_index 的 x,y 位置我能用 np.where 做吗

def radio_index2xy(radio_index,n_x,n_y): 

    con1 = radio_index <= n_x

    con2 = (radio_index > n_x) & (radio_index <= n_x+n_y-1)

    con3 = (radio_index > n_x+n_y-1) & (radio_index <= 2*n_x+n_y-2)

    con4 = (radio_index > 2*n_x+n_y-2) & (radio_index <= 2*n_x+2*n_y-4)

    condlist = [[con1],[con2],[con3],[con4]]

    choicelist = [[x_pos = radio_index -1 ,y_pos = 0],\

                  [(x_pos = n_x -1),(y_pos = radio_index - n_x)],\

                  [(x_pos = (n_x-1)-(radio_index-n_x-n_y+1)),(y_pos = n_y -1)],\

                  [（x_pos = 0),(y_pos = 2*n_x+2*n_y-4-radio_index+1)]]

    np.select(condlist,choicelist)

    return x_pos,y_pos 

if radio_index <= n_x:

    x_pos = radio_index -1

    y_pos = 0

elif radio_index > n_x and radio_index <= n_x+n_y-1:

    x_pos = n_x -1

    y_pos = radio_index - n_x

elif radio_index > n_x+n_y-1 and radio_index <= 2*n_x+n_y-2:

    x_pos = (n_x-1)-(radio_index-n_x-n_y+1)

    y_pos = n_y -1

elif radio_index > 2*n_x+n_y-2 and radio_index <= 2*n_x+2*n_y-4:

    x_pos = 0

    y_pos = 2*n_x+2*n_y-4-radio_index+1

Answer 1

np.searchsorted 对于这种逻辑很有用：def radio_index2xy_v(radio_index, n_x, n_y):&nbsp; &nbsp; sgn = np.array([0, 1, 1, -1, -1, 0])&nbsp; &nbsp; col = np.array([-1, 1, 0, 1, 0, -1])&nbsp; &nbsp; coeffs = np.array([[-1, -1],&nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp;[0, -1],&nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp;[-n_x, n_x - 1],&nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp;[n_y - 1, 2*n_x + n_y - 2],&nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp;[2*n_x + 2*n_y - 3, 0],&nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp;[-1, -1]])&nbsp; &nbsp; cusps = np.cumsum([0, n_x, n_y-1, n_x-1, n_y-2])&nbsp; &nbsp; idx = cusps.searchsorted(radio_index)&nbsp; &nbsp; out = coeffs[idx]&nbsp; &nbsp; out[np.arange(idx.size), col[idx]] += sgn[idx] * radio_index&nbsp; &nbsp; return out演示：>>> radio_index2xy_v(np.arange(20), 5, 4)array([[-1, -1],&nbsp; &nbsp; &nbsp; &nbsp;[ 0,&nbsp; 0],&nbsp; &nbsp; &nbsp; &nbsp;[ 0,&nbsp; 1],&nbsp; &nbsp; &nbsp; &nbsp;[ 0,&nbsp; 2],&nbsp; &nbsp; &nbsp; &nbsp;[ 0,&nbsp; 3],&nbsp; &nbsp; &nbsp; &nbsp;[ 0,&nbsp; 4],&nbsp; &nbsp; &nbsp; &nbsp;[ 1,&nbsp; 4],&nbsp; &nbsp; &nbsp; &nbsp;[ 2,&nbsp; 4],&nbsp; &nbsp; &nbsp; &nbsp;[ 3,&nbsp; 4],&nbsp; &nbsp; &nbsp; &nbsp;[ 3,&nbsp; 3],&nbsp; &nbsp; &nbsp; &nbsp;[ 3,&nbsp; 2],&nbsp; &nbsp; &nbsp; &nbsp;[ 3,&nbsp; 1],&nbsp; &nbsp; &nbsp; &nbsp;[ 3,&nbsp; 0],&nbsp; &nbsp; &nbsp; &nbsp;[ 2,&nbsp; 0],&nbsp; &nbsp; &nbsp; &nbsp;[ 1,&nbsp; 0],&nbsp; &nbsp; &nbsp; &nbsp;[-1, -1],&nbsp; &nbsp; &nbsp; &nbsp;[-1, -1],&nbsp; &nbsp; &nbsp; &nbsp;[-1, -1],&nbsp; &nbsp; &nbsp; &nbsp;[-1, -1],&nbsp; &nbsp; &nbsp; &nbsp;[-1, -1]])