n dimensional grid in Python / numpy

前端 未结 5 706
忘掉有多难
忘掉有多难 2021-01-23 06:57

I have an unknown number n of variables that can range from 0 to 1 with some known step s, with the condition that they sum up to 1. I want to create a

5条回答
  •  故里飘歌
    2021-01-23 07:37

    We can think of this as a problem of dividing some fixed number of things (1/s in this case and represented using sum_left parameter) between some given number of bins (n in this case). The most efficient way I can think of doing this is using a recursion:

    In [31]: arr = []
    In [32]: def fun(n, sum_left, arr_till_now):
        ...:     if n==1:
        ...:         n_arr = list(arr_till_now)
        ...:         n_arr.append(sum_left)
        ...:         arr.append(n_arr)
        ...:     else:
        ...:         for i in range(sum_left+1):
        ...:             n_arr = list(arr_till_now)
        ...:             n_arr.append(i)
        ...:             fun(n-1, sum_left-i, n_arr)
    

    This would give an output like:

    In [36]: fun(n, n, [])
    In [37]: arr
    Out[37]: 
    [[0, 0, 3],
     [0, 1, 2],
     [0, 2, 1],
     [0, 3, 0],
     [1, 0, 2],
     [1, 1, 1],
     [1, 2, 0],
     [2, 0, 1],
     [2, 1, 0],
     [3, 0, 0]]
    

    And now I can convert it to a numpy array to do an elementwise multiplication:

    In [39]: s = 0.33
    In [40]: arr_np = np.array(arr)
    In [41]: arr_np * s
    Out[41]: 
    array([[ 0.        ,  0.        ,  0.99999999],
           [ 0.        ,  0.33333333,  0.66666666],
           [ 0.        ,  0.66666666,  0.33333333],
           [ 0.        ,  0.99999999,  0.        ],
           [ 0.33333333,  0.        ,  0.66666666],
           [ 0.33333333,  0.33333333,  0.33333333],
           [ 0.33333333,  0.66666666,  0.        ],
           [ 0.66666666,  0.        ,  0.33333333],
           [ 0.66666666,  0.33333333,  0.        ],
           [ 0.99999999,  0.        ,  0.        ]])
    

提交回复
热议问题