Find matching pairs and lay them out as columns in polars

Say I have this:

df = polars.DataFrame(dict(
  j=numpy.random.randint(10, 99, 9),
  k=numpy.tile([1, 2, 2], 3),
  ))
  
 j (i64)  k (i64)
 47       1
 22       2
 82       2
 19       1
 85       2
 15       2
 89       1
 74       2
 26       2
shape: (9, 2)

where column k is kind of a marker - 1 starts and then there are one or more 2s (in the above example always two for simplicity, but in practice one or more). I'd like to get values in j that correspond to k=1 and the last corresponding k=2. For the above:

 j (i64)  k (i64)
 47       1 >-\
 22       2   | these are the 1 and the last of its matching 2s
 82       2 <-/
 19       1 >-\
 85       2   | these are the 1 and the last of its matching 2s
 15       2 <-/
 89       1 >-\
 74       2   | these are the 1 and the last of its matching 2s
 26       2 <-/
shape: (9, 2)

and I'd like to put these in two columns, so I get this:

 j (i64)  k (i64)
 47       82
 19       15
 89       26
shape: (9, 2)

How would I approach this in polars?

You can filter simply by looking for k=1 or when the next k, e.g. a shift, is 1:

df.select(
    j=pl.col('j').filter(pl.col('k') == 1),
    k=pl.col('j').filter(pl.col('k').shift(-1).fill_null(1) == 1),
)

shape: (3, 2)
┌─────┬─────┐
│ j   ┆ k   │
│ --- ┆ --- │
│ i32 ┆ i32 │
╞═════╪═════╡
│ 47  ┆ 82  │
│ 19  ┆ 15  │
│ 89  ┆ 26  │
└─────┴─────┘

import polars
import numpy

def construct_example(seed, n):
    numpy.random.seed(seed)
    ks = []
    js = []
    expected_res = []
    for i in range(n):
        ntwos = numpy.random.randint(1, 4)
        ks.extend([1] + [2 for j in range(ntwos)])
        ijs = numpy.random.randint(10, 99, ntwos + 1)
        js.extend(list(ijs))
        expected_res.append((ijs[0], ijs[-1]))
    df = polars.DataFrame(dict(j=js, k=ks))
    return df, expected_res

def solve(df):
    jarr = list(df['j']) + [None]
    karr = list(df['k']) + [1]
    res = []
    for i, (j, k) in enumerate(zip(jarr, karr)):
        if k == 1 and j is not None:
            res.append((j, jarr[i+karr[i+1:].index(1)]))
    return res

df, expected_res = construct_example(42, 10)

assert solve(df) == expected_res

print(list(df.iter_rows()))
print(expected_res)

prints

[(61, 1), (24, 2), (81, 2), (70, 2), (92, 1), (96, 2), (84, 1), (97, 2), (33, 2), (12, 2), (62, 1), (11, 2), (97, 2), (47, 1), (11, 2), (73, 2), (42, 1), (85, 2), (31, 1), (98, 2), (58, 2), (68, 1), (51, 2), (69, 2), (89, 2), (71, 1), (71, 2), (56, 2), (71, 2), (64, 1), (73, 2), (12, 2), (60, 2)]
[(61, 70), (92, 96), (84, 12), (62, 97), (47, 73), (42, 85), (31, 58), (68, 89), (71, 71), (64, 60)]

Explanation:

The function construct_example creates example data for n row groups, where the number of 2's per row groups can vary, and returns the corresponding polars.DataFrame and the expected pairs expected_res (as a list of tuples).

The function solve takes any such dataframe (assuming it satisfies the conditions of having only the said two markers, no two consective 1's and ending in a 2) and computes the matches as follows:

Add an extra row with k=1 and j=None, then iterate through the rows (indexed by i) and whenever you encounter a k=1 and corresponding j that is not None, take the j as the first element and then find the index of the next 1 (equal to i+1 plus the index of the first 1 when considering only values below/after), hence the corresponding j for the second element must sit at index i+karr[i+1:].index(1).