Get a list of rows starting from the same value as current row in pandas dataframe

Question

I have a dataframe that I'd like to expand with a new column which would contain/match the list of all ids if they fully contain the row string_value

id  string_value
1   The quick brown fox 
2   The quick brown fox jumps  
3   The quick brown fox jumps over 
4   The quick brown fox jumps over the lazy dog
5   The slow 
6   The slow brown fox

Desired output

id  string_value                                new_columns
1   The quick brown fox                         [2, 3, 4]
2   The quick brown fox jumps                   [3, 4]
3   The quick brown fox jumps over              [4]
4   The quick brown fox jumps over the lazy dog []
5   The slow                                    [6]
6   The slow brown fox                          []

Thanks

mozway · Answer 1 · 2022-09-09T13:39:54.927

2

You can't easily vectorize this, but you can use a custom function:

def accumulate(s):
    ref = None
    prev = s.index[0]
    out = {}
    for i, val in s.items():
        if ref and val.startswith(ref):
            tmp.append(prev)
        else:
            tmp = []
        ref = val
        prev = i
        out[i] = tmp.copy()

    # invert dictionary
    out2 = {}
    for v,l in out.items():
        for k in l:
            out2.setdefault(k, []).append(v)
    
    return pd.Series(out2)

df['new_columns'] = df['id'].map(accumulate(df.set_index('id')['string_value'].sort_values()))

output:

   id                                 string_value new_columns
0   1                          The quick brown fox   [2, 3, 4]
1   2                    The quick brown fox jumps      [3, 4]
2   3               The quick brown fox jumps over         [4]
3   4  The quick brown fox jumps over the lazy dog         NaN
4   5                                     The slow         [6]
5   6                           The slow brown fox         NaN

empty lists

to have empty lists in the output in place of NaN, change the "invert dictionary" code to:

    # invert dictionary
    out2 = {i: [] for i in s.index}
    for v,l in out.items():
        for k in l:
            out2[k].append(v)

edited Sep 09 '22 at 13:39

answered Sep 09 '22 at 13:07

mozway

194,879
13
39
75

Thanks. I get a slightly different results in 'new column' – Nadiia Sep 09 '22 at 13:36
What do you mean? The empty lists? – mozway Sep 09 '22 at 13:37
print(df.new_columns.values) produces [list([2]) nan list([4]) nan list([6]) nan] for the same df for me – Nadiia Sep 09 '22 at 14:50
Can you provide a DataFrame constructor for reproducibility? – mozway Sep 09 '22 at 14:50
sorry. figured it out. extra spaces at the end. – Nadiia Sep 09 '22 at 14:53
That's what I was betting on ;) – mozway Sep 09 '22 at 14:53

score 2 · Answer 2 · answered Sep 09 '22 at 16:15

Another way to do it using numpy:

# Check if a string starts with any other strings
mask = np.vstack([df["string_value"].str.contains(s) for s in df["string_value"]])

# We don't want to match a string to itself
# so mask out the diagonal
np.fill_diagonal(mask, False)

# Result
ids = df["id"].to_numpy()
df["new_column"] = [ids[m] for m in mask]

score 1 · Accepted Answer · answered Sep 09 '22 at 16:04

Here's another custom function you can consider. Assuming df is this:

   id                                 string_value
0   1                          The quick brown fox
1   2                    The quick brown fox jumps
2   3               The quick brown fox jumps over
3   4  The quick brown fox jumps over the lazy dog
4   5                                     The slow
5   6                           The slow brown fox

The custom function is

def match_string(string_value):
    idx_list = []
    for idx, strg in list(zip(df['id'], df['string_value'])):
        if strg == string_value:
            continue
        if string_value in strg:
            idx_list.append(idx)
    return idx_list

Then use lambda function:

df['new_columns'] = df['string_value'].apply(lambda x: match_string(x))
print(df)

   id                                 string_value new_columns
0   1                          The quick brown fox   [2, 3, 4]
1   2                    The quick brown fox jumps      [3, 4]
2   3               The quick brown fox jumps over         [4]
3   4  The quick brown fox jumps over the lazy dog          []
4   5                                     The slow         [6]
5   6                           The slow brown fox          []

Get a list of rows starting from the same value as current row in pandas dataframe

3 Answers3

empty lists