Find the next number value after a sublist in a list-CodePudding

I have two lists called x and y.

x = [[1, 2, 3], [4, 5, 6], [7, 8, 9]]

y = [4, 5]

I am trying to find the list that starts with y and get the following item (6 in this case). Also, i want to do this with a custom function. It must work like below

my_func(y)

takes the y and returns the number that comes next and equalize a variable to the number. Like, variable = 6

CodePudding user response：

Although you can do this with lists and loops, it is probably quickest to do it with numpy arrays:

import numpy as np
def get_next_value(x, y):
    x = np.array(x)
    y = np.array(y)
    num_to_check = y.shape[0]
    row_to_get_value = np.where((x[:,0:num_to_check] == y).sum(axis = 1) == num_to_check)
    desired_value = x[row_to_get_value, num_to_check]
         
    # Handle error if no match is found
    if desired_value.size > 0:
        return desired_value.item()
    else:
        return None

x = [[1, 2, 3], [4, 5, 6], [7, 8, 9]]
y = [4, 5]
get_next_value(x, y)
# 6

CodePudding user response：

IIUC, this should work for you:

x = [[1, 2, 3], [4, 5, 6], [7, 8, 9]]
y = [4, 5]

def my_func(y):
    output, ly = [], len(y)
    for l in x: 
        if l[:ly] == y: output.append(l[ly])
    return output

my_func(y)

Output:

[6]