你可以在numpy数组上广播字典定义吗?

时间:2013-03-25 02:36:04

标签: python numpy

我正在研究使用numpy的随机游走模拟的一个非常简单的例子。我的教授坚持认为我们尽可能多地使用numpy的广播功能而不是for循环,我想知道是否可以广播字典定义。

e.g。我有阵列[E W N S]。使用字典遍历该数组将导致[[1,0] [-1,0] [0,1] [0,-1]]。

import numpy as np
import matplotlib.pyplot as plt

def random_path(origin, nsteps, choices, choice_probs, choice_map):
    directions = np.random.choice(choices, size=(15,), p=choice_probs)
    print directions

def main():
    directions = ['N', 'S', 'E', 'W']
    dir_probabilities = [.2, .3, .45, .05]
    dir_map = {'N': [0, 1], 'S': [0, -1], 'E': [1, 0], 'W': [-1, 0]}
    origin = [0, 0]

    np.random.seed(12345)
    path = random_path(origin, 15, directions, dir_probabilities, dir_map)

main()

1 个答案:

答案 0 :(得分:2)

为什么不忽略实际的方向标签,只是将方向存储为(4,2)形状的numpy数组?然后你就可以直接索引到那个数组。

def random_path(origin, nsteps, choices, choice_probs, choice_map):
    directions = np.random.choice(choices, size=(15,), p=choice_probs)
    return directions

dir_map = np.array([[0,1], [0,-1], [1,0], [-1,0]])
# Everything else is the same as defined by OP

path_directions = random_path(origin, 15, np.arange(4), dir_probabilities, dir_map)
path = dir_map[path_directions]

现在path是一个(15,2)形的numpy数组,包含来自dir_map的移动序列。