Search code examples
pythonarrayspandasdataframeinsert

Modyfing 2d array


I'm stuck on trying to modify 2d array... Nothing I try seem to work... I'm trying to write a function that will add a value to its specific location in the numbers column...

import pandas as pd

def twod_array(num):
    data = {"group": [-1, 0, 1, 2],
            'numbers': [[2], [14, 15], [16, 17], [19, 20, 21]],
            }
    df = pd.DataFrame(data=data)
    print(df)
    return 0

Currently it prints this:

   group       numbers
0     -1           [2]
1      0      [14, 15]
2      1      [16, 17]
3      2  [19, 20, 21]

What I'd like to do is to add a value based on the passed input, so for example if I pass 14.5 as a num, this is the output I'd like to see:

   group       numbers
0     -1           [2]
1      0      [14,14.5 15]
2      1      [16, 17]
3      2  [19, 20, 21]

Another example: If I pass 18 as a num:

   group       numbers
0     -1           [2]
1      0      [14, 15]
2      1      [16, 17, 18]
3      2  [19, 20, 21]

I'm hoping someone can help with this.


Solution

  • df = pd.DataFrame({"group": [-1, 0, 1, 2],
            'numbers': [[2], [14, 15], [16, 17], [19, 20, 21]],
            })
    arr = df['numbers'].to_list()
    in_num = 18
    
    for i, sub_arr in enumerate(arr):
          for j, num in enumerate(sub_arr):
                if arr[i][j]>in_num:
                      if j!=0: arr[i].insert(j,in_num)
    
                      else: arr[i-1].insert(-1 ,in_num)
    
    df['numbers'] = arr