Home > Software engineering >  How to recursively extract values from a pandas DataFrame?
How to recursively extract values from a pandas DataFrame?

Time:11-15

I have the following pandas DataFrame:

df = pd.DataFrame([
  [3, 2, 5, 2],
  [8, 5, 4, 2],
  [9, 0, 8, 6],
  [9, 2, 7, 1],
  [1, 9, 2, 3],
  [8, 1, 1, 6],
  [8, 8, 0, 0],
  [0, 1, 3, 0],
  [2, 4, 5, 3],
  [4, 0, 9, 7]
])

I am trying to write a recursive function that extracts all the possible paths up until 3 iterations:

enter image description here

and saves them into a list. Several attempts but no results to post.

Desired Output:

[
    [0, 3, 9, 4],
    [0, 3, 9, 0],
    [0, 3, 9, 9],
    [0, 3, 9, 7],
    [0, 3, 2, 9],
    [0, 3, 2, 0],
    ...
]

Represented as a tree, this is how it looks like:

enter image description here

CodePudding user response:

Since you use numeric naming for both rows and columns in your dataframe, it's faster to convert the frame to a 2-D numpy array. Try this;

arr = df.to_numpy()

staging = [[0]]
result = []
while len(staging) > 0:
    s = staging.pop(0)
    if len(s) == 4:
        result.append(s)
    else:
        i = s[-1]
        for j in range(4):
            staging.append(s   [arr[i, j]])
  • Related