Rebuild pandas Dataframe with if Statements

As allready asked in another Question with the Title "Rebuild pandas Dataframe" i still have some Questions about going on with even more columns.

Situation: I have a Dataframe with 4 Columns, the Values inside the Columns are pretty Random. like this example:

df = pd.DataFrame({'col1': ['id 1', 'id 2', 'test 3', 'test 4'],
           'col2': ['test 1', 'test 2',
                    'ne 5261', 'id 4'],
           'col3': ['Number 12344', 'Number 21612','id 3','Number 1131'],
           'col4':['ne 315','Number 1264777','ne 1415','ne 52']})

My Goal is to have a Dataframe in wich each Col has only the Values begining with the Same substring like this example:

What allready worked for 3 columns is the following code (from The last Question credits to: @AndrejKesely):

def key_fn(x):
if 'id' in x:
    return 0
if 'test' in x:
    return 1
if 'Number' in x:
    return 2
return 3
df = pd.DataFrame([sorted(l, key=key_fn) for l in df.values], columns=df.columns)

print(df)

Since i now have 4 cols i added another if Statement to the Function, looking as follows:

def key_fn(x):
if 'id' in x:
    return 0
if 'test' in x:
    return 1
if 'Number' in x:
    return 2
if 'ne' in x:
    return 3
return 4
df = pd.DataFrame([sorted(l, key=key_fn) for l in df.values], columns=df.columns)

This gives me following output:

This is a small example, when i understand how it works i need to apply it to a Total of 17 Columns. Thank you in advance for your Help!