Создание массива массивов занимает слишком много времени

Question

Я пытаюсь создать массив всех возможных массивов вида

[x,y,z,(n-x-y-z),-20*x+0*y+-70*z+-20*(n-x-y-z)]

для данного n, где x, y и z являются целыми числами.

Он делает это правильно, но n=20 занимает около 2,3 секунд, а n=100 кажется, что это вечно.

Я знаю, что длина этого массива массивов увеличивается факториально с n.

Насколько я понимаю, numpy может позволить более эффективно создавать такой список, но я новичок в Python. Любые предложения о том, как сделать эту задачу более эффективно?

def total_outcomes(n):
  return [[x,y,z,(n-x-y-z),-20*x+0*y+-70*z+-20*(n-x-y-z)] for x in range(0,n+1) for y in range(0,n-x+1) for z in range(0,n-x-y+1)]

Answer 1

Если я правильно понимаю, вы хотите перебирать такие комбинации x, y и z, чтобы их сумма не превышала n.Если это так, вы можете создать все комбинации x, y и z от 0 до n, а затем создать логическую маску, сообщающую, что их сумма не больше n, и извлечь правильные комбинации.Проверьте, если это то, что вы хотите:

import numpy as np

def total_outcomes(n):
    x = np.arange(n+1)
    y = np.arange(n+1)
    z = np.arange(n+1)
    xyz_combinations = np.array(np.meshgrid(x, y, z)).T.reshape(-1,3)
    sums = np.sum(xyz_combinations, axis=1)
    mask = sums <= n
    xyz_combinations = xyz_combinations[mask]
    x = (xyz_combinations[:,0]).reshape(-1,1)
    y = (xyz_combinations[:,1]).reshape(-1,1)
    z = (xyz_combinations[:,2]).reshape(-1,1)
    c4 = (n-x-y-z).reshape(-1,1)
    #c5 = value_argmaxrcev_bias_scrutiny(n)*x+value_argmaxrcev_bias_noscrutiny(n)*y+value_argmaxrcev_nobias_scrutiny(n)*z+value_argmaxrcev_nobias_noscrutiny(n)*(c4)
    c5 = ((-20)*x+0*y+(-70)*z+(-20)*(c4)).reshape(-1,1)
    return np.hstack((x,y,z,c4,c5))

Он работает довольно быстро;в моем случае вычисление для n = 100 выполняется мгновенно.Время вычислений также растет довольно быстро;при n = 1000 он застревает на ~ 100 секунд.Как уже упоминалось в комментарии, вы не сможете преодолеть вычислительную сложность этой задачи.Приведение проблемы к numpy на самом деле просто делает вычисления более низкоуровневыми (поскольку numpy написано на C и упаковано в python), и, следовательно, быстрее, чем зацикливание на python.

Этот код для n= 5, генерирует:

array([[   0,    0,    0,    5, -100],
       [   0,    1,    0,    4,  -80],
       [   0,    2,    0,    3,  -60],
       [   0,    3,    0,    2,  -40],
       [   0,    4,    0,    1,  -20],
       [   0,    5,    0,    0,    0],
       [   1,    0,    0,    4, -100],
       [   1,    1,    0,    3,  -80],
       [   1,    2,    0,    2,  -60],
       [   1,    3,    0,    1,  -40],
       [   1,    4,    0,    0,  -20],
       [   2,    0,    0,    3, -100],
       [   2,    1,    0,    2,  -80],
       [   2,    2,    0,    1,  -60],
       [   2,    3,    0,    0,  -40],
       [   3,    0,    0,    2, -100],
       [   3,    1,    0,    1,  -80],
       [   3,    2,    0,    0,  -60],
       [   4,    0,    0,    1, -100],
       [   4,    1,    0,    0,  -80],
       [   5,    0,    0,    0, -100],
       [   0,    0,    1,    4, -150],
       [   0,    1,    1,    3, -130],
       [   0,    2,    1,    2, -110],
       [   0,    3,    1,    1,  -90],
       [   0,    4,    1,    0,  -70],
       [   1,    0,    1,    3, -150],
       [   1,    1,    1,    2, -130],
       [   1,    2,    1,    1, -110],
       [   1,    3,    1,    0,  -90],
       [   2,    0,    1,    2, -150],
       [   2,    1,    1,    1, -130],
       [   2,    2,    1,    0, -110],
       [   3,    0,    1,    1, -150],
       [   3,    1,    1,    0, -130],
       [   4,    0,    1,    0, -150],
       [   0,    0,    2,    3, -200],
       [   0,    1,    2,    2, -180],
       [   0,    2,    2,    1, -160],
       [   0,    3,    2,    0, -140],
       [   1,    0,    2,    2, -200],
       [   1,    1,    2,    1, -180],
       [   1,    2,    2,    0, -160],
       [   2,    0,    2,    1, -200],
       [   2,    1,    2,    0, -180],
       [   3,    0,    2,    0, -200],
       [   0,    0,    3,    2, -250],
       [   0,    1,    3,    1, -230],
       [   0,    2,    3,    0, -210],
       [   1,    0,    3,    1, -250],
       [   1,    1,    3,    0, -230],
       [   2,    0,    3,    0, -250],
       [   0,    0,    4,    1, -300],
       [   0,    1,    4,    0, -280],
       [   1,    0,    4,    0, -300],
       [   0,    0,    5,    0, -350]])