如何用单个numpy数组操作替换这个三重For循环?

时间:2018-01-27 05:39:31

标签: python arrays numpy graph-theory

对于基于回合的游戏,我想计算每个玩家可以在地图的每个区域上移动或产生的最大单位数。

我需要计算的所有数据都存储在几个numpy数组中,但我很难找到高级数组索引技术来尽可能快地进行计算。

为了帮助解决这个问题,我用一些For循环以最简单的方式重写了这个函数:

import numpy as np

def get_max_units_on_zone_per_player(unitCountPerPlayer, zoneOwner, playerAvailableUnits, zoneLinks, blockedMovesPerPlayer):
    """
    Parameters
    ----------
    unitCountPerPlayer: np.array((zoneCount, playerCount), dtype=int)
        How many units each player has on a zone
    zoneOwner: np.array(zoneCount, dtype=int)
        Which player is owning a zone (-1 for none)
    playerAvailableUnits: np.array(playerCount, dtype=int)
        How many units each player can spawn
    zoneLinks: np.array((zoneCount, zoneCount), dtype=int)
        > 0 if zone1 is connected to zone2 (directed and weighted graph)
    blockedMovesPerPlayer: np.array((playerCount, zoneCount, zoneCount), dtype=bool)
        True if player can not move from zone1 to zone2

    Returns
    -------
    np.array((zoneCount, playerCount), dtype=int)
        Maximum count of units each player can have on each zone
    """

    zoneCount, playerCount = unitCountPerPlayer.shape

    # Adding units already on zone
    result = np.zeros((zoneCount, playerCount), dtype=int) + unitCountPerPlayer

    for p in xrange(playerCount):
        for z1 in xrange(zoneCount):

            if zoneOwner[z1] in (-1, p):
                # Player can spawn on neutral or owned zones
                result[z1, p] += playerAvailableUnits[p]

            for z2 in xrange(zoneCount):

                if zoneLinks[z1, z2] > 0 and not blockedMovesPerPlayer[p, z1, z2]:
                    # If z1 and z2 are connected and player can move from z1 to z2, adding units count on z1 to z2
                    result[z2, p] += unitCountPerPlayer[z1, p]
    return result

问题是我无法使用此功能,每次通话大约需要30毫秒,我确信它可以重写为一些小块操作,处理时间不到5毫秒。

有人可以帮我这个吗?还有一步一步的过程,以便下次我可以自己做吗?我已经多次阅读关于数组和索引的numpy文档,但它远非一尘不染,我无法弄清楚。

编辑:根据要求,以下是一些可用作样本的随机数据:

zoneCount=8 ; playerCount=2

unitCountPerPlayer:
[[1 2]
 [1 3]
 [1 3]
 [3 2]
 [1 2]
 [3 2]
 [0 2]
 [3 2]]

zoneOwner:
[ 1  0 -1 -1 -1  0 -1 -1]

playerAvailableUnits:
[2 2]

zoneLinks:
[[0 1 1 1 0 1 0 0]
 [1 0 0 1 0 0 0 1]
 [1 1 1 1 0 1 0 1]
 [0 1 1 1 1 0 1 0]
 [0 0 1 1 1 0 1 1]
 [0 0 1 1 1 1 1 1]
 [1 0 0 0 0 1 0 1]
 [1 1 1 1 0 1 1 1]]

blockedMovesPerPlayer:
[[[False False False False False False False False]
  [ True False False False False False False False]
  [ True False False False False False False False]
  [False False False False False False False False]
  [False False False False False False False False]
  [False False False False False False False False]
  [ True False False False False False False False]
  [ True False False False False False False False]]

 [[False  True False False False  True False False]
  [False False False False False False False False]
  [False  True False False False  True False False]
  [False  True False False False False False False]
  [False False False False False False False False]
  [False False False False False  True False False]
  [False False False False False  True False False]
  [False  True False False False  True False False]]]

get_max_units_on_zone_per_player():
[[ 1 14]
 [11  3]
 [15 18]
 [18 20]
 [10 10]
 [13  2]
 [12 12]
 [14 18]]

复制/粘贴数据:

zoneCount = 8
playerCount = 2

unitCountPerPlayer = np.array([[1,2], [1,3], [1,3], [3,2],
                               [1,2], [3,2], [0,2], [3,2]])

zoneOwner = np.array([1, 0, -1, -1, -1, 0, -1, -1])

playerAvailableUnits = np.array([2,2])

zoneLinks = np.array([[0,1,1,1,0,1,0,0], [1,0,0,1,0,0,0,1],
                      [1,1,1,1,0,1,0,1], [0,1,1,1,1,0,1,0],
                      [0,0,1,1,1,0,1,1], [0,0,1,1,1,1,1,1],
                      [1,0,0,0,0,1,0,1], [1,1,1,1,0,1,1,1]])

bmpp = [[[False, False, False, False, False, False, False, False],
         [ True, False, False, False, False, False, False, False],
         [ True, False, False, False, False, False, False, False],
         [False, False, False, False, False, False, False, False],
         [False, False, False, False, False, False, False, False],
         [False, False, False, False, False, False, False, False],
         [ True, False, False, False, False, False, False, False],
         [ True, False, False, False, False, False, False, False]],
        [[False,  True, False, False, False,  True, False, False],
         [False, False, False, False, False, False, False, False],
         [False,  True, False, False, False,  True, False, False],
         [False,  True, False, False, False, False, False, False],
         [False, False, False, False, False, False, False, False],
         [False, False, False, False, False,  True, False, False],
         [False, False, False, False, False,  True, False, False],
         [False,  True, False, False, False,  True, False, False]]]
blockedMovesPerPlayer = np.array(bmpp)

1 个答案:

答案 0 :(得分:1)

[更新:实施numpy方式,避免for循环播放

以下是get_max_units_on_zone_per_player()实现:

def get_max_units_on_zone_per_player(unitCountPerPlayer, zoneOwner, playerAvailableUnits, zoneLinks, blockedMovesPerPlayer):
    result = unitCountPerPlayer.copy()
    result[zoneOwner < 0] += playerAvailableUnits
    _z1 = np.where(zoneOwner >= 0)
    result[_z1, zoneOwner[_z1]] += playerAvailableUnits[zoneOwner[_z1]]
    _p, _z1, _z2 = np.where(np.logical_and(zoneLinks > 0, np.logical_not(blockedMovesPerPlayer)))
    np.add.at(result, [_z2, _p], unitCountPerPlayer[_z1, _p])
    return result

我使用以下设置测试了这两个实现:

zoneCount = 100
playerCount = 1000
maxUnits = 500

unitCountPerPlayer = np.random.randint(0, maxUnits, size=(zoneCount, playerCount))
zoneOwner = np.random.randint(-1, playerCount, size=zoneCount)
playerAvailableUnits = np.random.randint(0, maxUnits, size=playerCount)
zoneLinks = np.random.randint(0, maxUnits, size=(zoneCount, zoneCount))
blockedMovesPerPlayer = np.random.randint(0, 2, size=(playerCount, zoneCount, zoneCount), dtype=bool)

以下是测试结果(%timeit

  • fbparis的原始实现:

    每回路7.27 s±10 ms(平均值±标准偏差,7次运行,每次1次循环)

  • 我的实施:

    每循环645 ms±490μs(平均值±标准偏差,7次运行,每次循环1次)