def g(arr):
    arr[:, 0] = np.where(np.isnan(arr[:, 0]), arr[:, 0], np.round(arr[:, 0], 2))
    return arr

arr = g(arr.copy())

import numpy as np
import copy

def generate_test_case(test_case_id):
    def generate_ans(data):
        arr = data
        arr[:, 0] = np.where(np.isnan(arr[:, 0]), arr[:, 0], np.round(arr[:, 0], 2))
        return arr

    def define_test_input(test_case_id):
        if test_case_id == 1:
            arr = np.array([
                [0.21, 0.3212],
                [0.01, 0.61237],
                [0.66123, 0.03],
                [0.21, 0.18],
                [np.nan, 0.18],
            ])
        if test_case_id == 2:
            arr = np.array([
                [0.215, 0.3212],
                [0.01, 0.11237],
                [0.36123, 0.03],
                [0.21, 0.18],
                [np.nan, 0.18],
            ])
        return arr

    test_input = define_test_input(test_case_id)
    expected_result = generate_ans(copy.deepcopy(test_input))
    return test_input, expected_result

def exec_test(result, ans):
    try:
        np.testing.assert_array_almost_equal(result, ans, decimal=2)
        return 1
    except:
        return 0

exec_context = r"""
import numpy as np
arr = test_input
[insert]
result = arr
"""

def test_execution(solution: str):
    code = exec_context.replace("[insert]", solution)
    for i in range(2):
        test_input, expected_result = generate_test_case(i + 1)
        test_env = {"test_input": test_input}
        exec(code, test_env)
        assert exec_test(test_env["result"], expected_result)

import pandas as pd

df = pd.DataFrame([(.21, .3212), (.01, .61237), (.66123, .03), (.21, .18),(pd.NA, .18)],
                  columns=['dogs', 'cats'])

df = ... # put solution in this variable

def g(df):
    df['dogs'] = df['dogs'].apply(lambda x: round(x,2) if str(x) != '<NA>' else x)
    return df

df = g(df.copy())

import pandas as pd
import numpy as np
import copy

def generate_test_case(test_case_id):
    def generate_ans(data):
        df = data
        df["dogs"] = df["dogs"].apply(lambda x: round(x, 2) if str(x) != "<NA>" else x)
        return df

    def define_test_input(test_case_id):
        if test_case_id == 1:
            df = pd.DataFrame(
                [
                    (0.21, 0.3212),
                    (0.01, 0.61237),
                    (0.66123, 0.03),
                    (0.21, 0.18),
                    (pd.NA, 0.18),
                ],
                columns=["dogs", "cats"],
            )
        if test_case_id == 2:
            df = pd.DataFrame(
                [
                    (0.215, 0.3212),
                    (0.01, 0.11237),
                    (0.36123, 0.03),
                    (0.21, 0.18),
                    (pd.NA, 0.18),
                ],
                columns=["dogs", "cats"],
            )
        return df

    test_input = define_test_input(test_case_id)
    expected_result = generate_ans(copy.deepcopy(test_input))
    return test_input, expected_result

def exec_test(result, ans):
    try:
        pd.testing.assert_frame_equal(result, ans, check_dtype=False)
        return 1
    except:
        return 0

exec_context = r"""
import pandas as pd
import numpy as np
df = test_input
[insert]
result = df
"""

def test_execution(solution: str):
    code = exec_context.replace("[insert]", solution)
    for i in range(2):
        test_input, expected_result = generate_test_case(i + 1)
        test_env = {"test_input": test_input}
        exec(code, test_env)
        assert exec_test(test_env["result"], expected_result)


def g(arr):
    for i in range(arr.shape[0]):
        if not np.isnan(arr[i, 0]) and not np.isnan(arr[i, 1]):
            arr[i, 0] = np.round(arr[i, 0], 2)
            arr[i, 1] = np.round(arr[i, 1], 2)
    return arr

arr = g(arr.copy())

import numpy as np
import copy

def generate_test_case(test_case_id):
    def generate_ans(data):
        arr = data
        for i in range(arr.shape[0]):
            if not np.isnan(arr[i, 0]) and not np.isnan(arr[i, 1]):
                arr[i, 0] = np.round(arr[i, 0], 2)
                arr[i, 1] = np.round(arr[i, 1], 2)
        return arr

    def define_test_input(test_case_id):
        if test_case_id == 1:
            arr = np.array([
                [0.21, 0.3212],
                [0.01, 0.61237],
                [0.66123, np.nan],
                [0.21, 0.18],
                [np.nan, 0.188],
            ])
        if test_case_id == 2:
            arr = np.array([
                [np.nan, 0.3212],
                [0.01, 0.61237],
                [0.66123, np.nan],
                [0.21, 0.18],
                [np.nan, 0.188],
            ])
        return arr

    test_input = define_test_input(test_case_id)
    expected_result = generate_ans(copy.deepcopy(test_input))
    return test_input, expected_result

def exec_test(result, ans):
    try:
        np.testing.assert_array_almost_equal(result, ans, decimal=2)
        return 1
    except:
        return 0

exec_context = r"""
import numpy as np
arr = test_input
[insert]
result = arr
"""

def test_execution(solution: str):
    code = exec_context.replace("[insert]", solution)
    for i in range(2):
        test_input, expected_result = generate_test_case(i + 1)
        test_env = {"test_input": test_input}
        exec(code, test_env)
        assert exec_test(test_env["result"], expected_result)

import pandas as pd

df = pd.DataFrame([(.21, .3212), (.01, .61237), (.66123, pd.NA), (.21, .18),(pd.NA, .188)],
                  columns=['dogs', 'cats'])

df = ... # put solution in this variable

def g(df):
    for i in df.index:
        if str(df.loc[i, 'dogs']) != '<NA>' and str(df.loc[i, 'cats']) != '<NA>':
            df.loc[i, 'dogs'] = round(df.loc[i, 'dogs'], 2)
            df.loc[i, 'cats'] = round(df.loc[i, 'cats'], 2)
    return df

df = g(df.copy())

import pandas as pd
import numpy as np
import copy

def generate_test_case(test_case_id):
    def generate_ans(data):
        df = data
        for i in df.index:
            if str(df.loc[i, "dogs"]) != "<NA>" and str(df.loc[i, "cats"]) != "<NA>":
                df.loc[i, "dogs"] = round(df.loc[i, "dogs"], 2)
                df.loc[i, "cats"] = round(df.loc[i, "cats"], 2)
        return df

    def define_test_input(test_case_id):
        if test_case_id == 1:
            df = pd.DataFrame(
                [
                    (0.21, 0.3212),
                    (0.01, 0.61237),
                    (0.66123, pd.NA),
                    (0.21, 0.18),
                    (pd.NA, 0.188),
                ],
                columns=["dogs", "cats"],
            )
        if test_case_id == 2:
            df = pd.DataFrame(
                [
                    (pd.NA, 0.3212),
                    (0.01, 0.61237),
                    (0.66123, pd.NA),
                    (0.21, 0.18),
                    (pd.NA, 0.188),
                ],
                columns=["dogs", "cats"],
            )
        return df

    test_input = define_test_input(test_case_id)
    expected_result = generate_ans(copy.deepcopy(test_input))
    return test_input, expected_result

def exec_test(result, ans):
    try:
        pd.testing.assert_frame_equal(result, ans, check_dtype=False)
        return 1
    except:
        return 0

exec_context = r"""
import pandas as pd
import numpy as np
df = test_input
[insert]
result = df
"""

def test_execution(solution: str):
    code = exec_context.replace("[insert]", solution)
    for i in range(2):
        test_input, expected_result = generate_test_case(i + 1)
        test_env = {"test_input": test_input}
        exec(code, test_env)
        assert exec_test(test_env["result"], expected_result)


import numpy as np

np.random.seed(10)
data = np.random.randint(1, 100, (10, 26))
list_of_my_columns = [0, 4, 25]

arr = ... # put solution in this variable

def g(arr, list_of_my_columns):
    avg_values = np.mean(arr[:, list_of_my_columns], axis=1)
    min_values = np.min(arr[:, list_of_my_columns], axis=1)
    max_values = np.max(arr[:, list_of_my_columns], axis=1)
    median_values = np.median(arr[:, list_of_my_columns], axis=1)
    return avg_values, min_values, max_values, median_values

avg_values, min_values, max_values, median_values = g(arr.copy(), list_of_my_columns.copy())

import numpy as np
import copy

def generate_test_case(test_case_id):
    def generate_ans(data):
        arr, list_of_my_columns = data
        avg_values = np.mean(arr[:, list_of_my_columns], axis=1)
        min_values = np.min(arr[:, list_of_my_columns], axis=1)
        max_values = np.max(arr[:, list_of_my_columns], axis=1)
        median_values = np.median(arr[:, list_of_my_columns], axis=1)
        return avg_values, min_values, max_values, median_values

    def define_test_input(test_case_id):
        if test_case_id == 1:
            np.random.seed(10)
            data = np.random.randint(1, 100, (10, 26))
            list_of_my_columns = [0, 4, 25]
        return data, list_of_my_columns

    test_input = define_test_input(test_case_id)
    expected_result = generate_ans(copy.deepcopy(test_input))
    return test_input, expected_result

def exec_test(result, ans):
    try:
        for r, a in zip(result, ans):
            np.testing.assert_array_almost_equal(r, a)
        return 1
    except:
        return 0

exec_context = r"""
import numpy as np
arr, list_of_my_columns = test_input
[insert]
result = avg_values, min_values, max_values, median_values
"""

def test_execution(solution: str):
    code = exec_context.replace("[insert]", solution)
    for i in range(1):
        test_input, expected_result = generate_test_case(i + 1)
        test_env = {"test_input": test_input}
        exec(code, test_env)
        assert exec_test(test_env["result"], expected_result)

import pandas as pd
import numpy as np

np.random.seed(10)
data = {}
for i in [chr(x) for x in range(65,91)]:
    data['Col '+i] = np.random.randint(1,100,10)
df = pd.DataFrame(data)
list_of_my_columns = ['Col A', 'Col E', 'Col Z']

df = ... # put solution in this variable

def g(df, list_of_my_columns):
    df['Avg'] = df[list_of_my_columns].mean(axis=1)
    df['Min'] = df[list_of_my_columns].min(axis=1)
    df['Max'] = df[list_of_my_columns].max(axis=1)
    df['Median'] = df[list_of_my_columns].median(axis=1)
    return df

df = g(df.copy(),list_of_my_columns.copy())

import pandas as pd
import numpy as np
import copy

def generate_test_case(test_case_id):
    def generate_ans(data):
        data = data
        df, list_of_my_columns = data
        df["Avg"] = df[list_of_my_columns].mean(axis=1)
        df["Min"] = df[list_of_my_columns].min(axis=1)
        df["Max"] = df[list_of_my_columns].max(axis=1)
        df["Median"] = df[list_of_my_columns].median(axis=1)
        return df

    def define_test_input(test_case_id):
        if test_case_id == 1:
            np.random.seed(10)
            data = {}
            for i in [chr(x) for x in range(65, 91)]:
                data["Col " + i] = np.random.randint(1, 100, 10)
            df = pd.DataFrame(data)
            list_of_my_columns = ["Col A", "Col E", "Col Z"]
        return df, list_of_my_columns

    test_input = define_test_input(test_case_id)
    expected_result = generate_ans(copy.deepcopy(test_input))
    return test_input, expected_result

def exec_test(result, ans):
    try:
        pd.testing.assert_frame_equal(result, ans, check_dtype=False)
        return 1
    except:
        return 0

exec_context = r"""
import pandas as pd
import numpy as np
df, list_of_my_columns = test_input
[insert]
result = df
"""

def test_execution(solution: str):
    code = exec_context.replace("[insert]", solution)
    for i in range(1):
        test_input, expected_result = generate_test_case(i + 1)
        test_env = {"test_input": test_input}
        exec(code, test_env)
        assert exec_test(test_env["result"], expected_result)


def g(arr):
    return np.sort(arr, order='time')

result = g(arr.copy())

import numpy as np
import copy

def generate_test_case(test_case_id):
    def generate_ans(data):
        arr = data
        return np.sort(arr, order='time')

    def define_test_input(test_case_id):
        if test_case_id == 1:
            arr = np.array([
                ('TGFb', 0.1, 2, -0.158406),
                ('TGFb', 1.0, 2, 0.039158),
                ('TGFb', 10.0, 2, -0.052608),
                ('TGFb', 0.1, 24, 0.157153),
                ('TGFb', 1.0, 24, 0.20603),
                ('TGFb', 10.0, 24, 0.13258),
                ('TGFb', 0.1, 48, -0.144209),
                ('TGFb', 1.0, 48, -0.09391),
                ('TGFb', 10.0, 48, -0.166819),
                ('TGFb', 0.1, 6, 0.097548),
                ('TGFb', 1.0, 6, 0.026664),
                ('TGFb', 10.0, 6, -0.008032)
            ], dtype=[('treatment', 'U4'), ('dose', 'f4'), ('time', 'i4'), ('VIM', 'f4')])
        return arr

    test_input = define_test_input(test_case_id)
    expected_result = generate_ans(copy.deepcopy(test_input))
    return test_input, expected_result

def exec_test(result, ans):
    try:
        np.testing.assert_array_equal(result, ans)
        return 1
    except:
        return 0

exec_context = r"""
import numpy as np
arr = test_input
[insert]
"""

def test_execution(solution: str):
    code = exec_context.replace("[insert]", solution)
    for i in range(1):
        test_input, expected_result = generate_test_case(i + 1)
        test_env = {"test_input": test_input}
        exec(code, test_env)
        assert exec_test(test_env["result"], expected_result)

def g(df):
    return df.sort_index(level='time')

result = g(df.copy())

import pandas as pd
import numpy as np
import copy

def generate_test_case(test_case_id):
    def generate_ans(data):
        df = data
        return df.sort_index(level="time")

    def define_test_input(test_case_id):
        if test_case_id == 1:
            df = pd.DataFrame(
                {
                    "VIM": [
                        -0.158406,
                        0.039158,
                        -0.052608,
                        0.157153,
                        0.206030,
                        0.132580,
                        -0.144209,
                        -0.093910,
                        -0.166819,
                        0.097548,
                        0.026664,
                        -0.008032,
                    ]
                },
                index=pd.MultiIndex.from_tuples(
                    [
                        ("TGFb", 0.1, 2),
                        ("TGFb", 1, 2),
                        ("TGFb", 10, 2),
                        ("TGFb", 0.1, 24),
                        ("TGFb", 1, 24),
                        ("TGFb", 10, 24),
                        ("TGFb", 0.1, 48),
                        ("TGFb", 1, 48),
                        ("TGFb", 10, 48),
                        ("TGFb", 0.1, 6),
                        ("TGFb", 1, 6),
                        ("TGFb", 10, 6),
                    ],
                    names=["treatment", "dose", "time"],
                ),
            )
        return df

    test_input = define_test_input(test_case_id)
    expected_result = generate_ans(copy.deepcopy(test_input))
    return test_input, expected_result

def exec_test(result, ans):
    try:
        pd.testing.assert_frame_equal(result, ans, check_dtype=False)
        return 1
    except:
        return 0

exec_context = r"""
import pandas as pd
import numpy as np
df = test_input
[insert]
"""

def test_execution(solution: str):
    code = exec_context.replace("[insert]", solution)
    for i in range(1):
        test_input, expected_result = generate_test_case(i + 1)
        test_env = {"test_input": test_input}
        exec(code, test_env)
        assert exec_test(test_env["result"], expected_result)


def g(arr):
    return np.sort(arr, order='VIM')

result = g(structured_array.copy())

import numpy as np
import copy

def generate_test_case(test_case_id):
    def generate_ans(data):
        arr = data
        return np.sort(arr, order='VIM')

    def define_test_input(test_case_id):
        if test_case_id == 1:
            arr = np.array([
                (('TGFb', 0.1, 2), -0.158406),
                (('TGFb', 1, 2), 0.039158),
                (('TGFb', 10, 2), -0.052608),
                (('TGFb', 0.1, 24), 0.157153),
                (('TGFb', 1, 24), 0.206030),
                (('TGFb', 10, 24), 0.132580),
                (('TGFb', 0.1, 48), -0.144209),
                (('TGFb', 1, 48), -0.093910),
                (('TGFb', 10, 48), -0.166819),
                (('TGFb', 0.1, 6), 0.097548),
                (('TGFb', 1, 6), 0.026664),
                (('TGFb', 10, 6), -0.008032)
            ], dtype=[('index', [('treatment', 'U4'), ('dose', 'f4'), ('time', 'i4')]), ('VIM', 'f4')])
        return arr

    test_input = define_test_input(test_case_id)
    expected_result = generate_ans(copy.deepcopy(test_input))
    return test_input, expected_result

def exec_test(result, ans):
    try:
        np.testing.assert_array_equal(result, ans)
        return 1
    except:
        return 0

exec_context = r"""
import numpy as np
structured_array = test_input
[insert]
"""

def test_execution(solution: str):
    code = exec_context.replace("[insert]", solution)
    for i in range(1):
        test_input, expected_result = generate_test_case(i + 1)
        test_env = {"test_input": test_input}
        exec(code, test_env)
        assert exec_test(test_env["result"], expected_result)

def g(df):
    return df.sort_values('VIM')

result = g(df.copy())

import pandas as pd
import numpy as np
import copy

def generate_test_case(test_case_id):
    def generate_ans(data):
        df = data
        return df.sort_values("VIM")

    def define_test_input(test_case_id):
        if test_case_id == 1:
            df = pd.DataFrame(
                {
                    "VIM": [
                        -0.158406,
                        0.039158,
                        -0.052608,
                        0.157153,
                        0.206030,
                        0.132580,
                        -0.144209,
                        -0.093910,
                        -0.166819,
                        0.097548,
                        0.026664,
                        -0.008032,
                    ]
                },
                index=pd.MultiIndex.from_tuples(
                    [
                        ("TGFb", 0.1, 2),
                        ("TGFb", 1, 2),
                        ("TGFb", 10, 2),
                        ("TGFb", 0.1, 24),
                        ("TGFb", 1, 24),
                        ("TGFb", 10, 24),
                        ("TGFb", 0.1, 48),
                        ("TGFb", 1, 48),
                        ("TGFb", 10, 48),
                        ("TGFb", 0.1, 6),
                        ("TGFb", 1, 6),
                        ("TGFb", 10, 6),
                    ],
                    names=["treatment", "dose", "time"],
                ),
            )
        return df

    test_input = define_test_input(test_case_id)
    expected_result = generate_ans(copy.deepcopy(test_input))
    return test_input, expected_result

def exec_test(result, ans):
    try:
        pd.testing.assert_frame_equal(result, ans, check_dtype=False)
        return 1
    except:
        return 0

exec_context = r"""
import pandas as pd
import numpy as np
df = test_input
[insert]
"""

def test_execution(solution: str):
    code = exec_context.replace("[insert]", solution)
    for i in range(1):
        test_input, expected_result = generate_test_case(i + 1)
        test_env = {"test_input": test_input}
        exec(code, test_env)
        assert exec_test(test_env["result"], expected_result)


def g(data):
    to_delete = ['2020-02-17', '2020-02-18']
    mask = np.isin(np.array([d[:10] for d in data['Date']]), to_delete, invert=True)
    return data[mask]

result = g(data.copy())

sp.head() with Date column including '2020-02-17' and '2020-02-18'

def g(df):
    to_delete = ['2020-02-17', '2020-02-18']
    return df[~(df.index.strftime('%Y-%m-%d').isin(to_delete))]

result = g(df.copy())


def g(data, dates):
    to_delete = ['2020-02-17', '2020-02-18']
    mask = ~np.isin(dates.astype('datetime64[D]'), np.array(to_delete, dtype='datetime64'))
    filtered_data = data[mask]
    formatted_dates = np.array([date.strftime('%d-%b-%Y %A') for date in dates[mask]])
    return np.core.records.fromarrays([formatted_dates] + [filtered_data[name] for name in data.dtype.names[1:]], names=data.dtype.names)

result = g(data, dates)

def g(df):
    to_delete = ['2020-02-17', '2020-02-18']
    df = df[~(df.index.strftime('%Y-%m-%d').isin(to_delete))]
    df.index = df.index.strftime('%d-%b-%Y %A')
    return df

result = g(df.copy())


import numpy as np

np.random.seed(10)
data = np.random.rand(10, 5)
corr = np.corrcoef(data, rowvar=False)

result = ... # put solution in this variable

def g(corr):
    mask = np.triu(np.ones(corr.shape), k=1).astype(bool)
    corr_triu = corr[mask]
    indices = np.argwhere(mask)
    filtered_indices = indices[corr_triu > 0.3]
    filtered_values = corr_triu[corr_triu > 0.3]
    return np.column_stack((filtered_indices, filtered_values))

result = g(corr.copy())

import numpy as np
import copy

def generate_test_case(test_case_id):
    def generate_ans(data):
        mask = np.triu(np.ones(data.shape), k=1).astype(bool)
        corr_triu = data[mask]
        indices = np.argwhere(mask)
        filtered_indices = indices[corr_triu > 0.3]
        filtered_values = corr_triu[corr_triu > 0.3]
        return np.column_stack((filtered_indices, filtered_values))

    def define_test_input(test_case_id):
        if test_case_id == 1:
            np.random.seed(10)
            data = np.random.rand(10, 5)
            corr = np.corrcoef(data, rowvar=False)
        return corr

    test_input = define_test_input(test_case_id)
    expected_result = generate_ans(copy.deepcopy(test_input))
    return test_input, expected_result

def exec_test(result, ans):
    try:
        np.testing.assert_array_almost_equal(result, ans)
        return 1
    except:
        return 0

exec_context = r"""
import numpy as np
corr = test_input
[insert]
"""

def test_execution(solution: str):
    code = exec_context.replace("[insert]", solution)
    for i in range(1):
        test_input, expected_result = generate_test_case(i + 1)
        test_env = {"test_input": test_input}
        exec(code, test_env)
        assert exec_test(test_env["result"], expected_result)

import pandas as pd
import numpy as np

np.random.seed(10)
df = pd.DataFrame(np.random.rand(10,5))
corr = df.corr()

result = ... # put solution in this variable

def g(corr):
    corr_triu = corr.where(~np.tril(np.ones(corr.shape)).astype(bool))
    corr_triu = corr_triu.stack()
    corr_triu.name = 'Pearson Correlation Coefficient'
    corr_triu.index.names = ['Col1', 'Col2']
    return corr_triu[corr_triu > 0.3].to_frame()

result = g(corr.copy())

import pandas as pd
import numpy as np
import copy

def generate_test_case(test_case_id):
    def generate_ans(data):
        corr = data
        corr_triu = corr.where(~np.tril(np.ones(corr.shape)).astype(bool))
        corr_triu = corr_triu.stack()
        corr_triu.name = "Pearson Correlation Coefficient"
        corr_triu.index.names = ["Col1", "Col2"]
        return corr_triu[corr_triu > 0.3].to_frame()

    def define_test_input(test_case_id):
        if test_case_id == 1:
            np.random.seed(10)
            df = pd.DataFrame(np.random.rand(10, 5))
            corr = df.corr()
        return corr

    test_input = define_test_input(test_case_id)
    expected_result = generate_ans(copy.deepcopy(test_input))
    return test_input, expected_result

def exec_test(result, ans):
    try:
        pd.testing.assert_frame_equal(result, ans, check_dtype=False)
        return 1
    except:
        return 0

exec_context = r"""
import pandas as pd
import numpy as np
corr = test_input
[insert]
"""

def test_execution(solution: str):
    code = exec_context.replace("[insert]", solution)
    for i in range(1):
        test_input, expected_result = generate_test_case(i + 1)
        test_env = {"test_input": test_input}
        exec(code, test_env)
        assert exec_test(test_env["result"], expected_result)


def g(corr):
    corr_triu = np.triu(corr, k=1)
    indices = np.where(corr_triu > 0.3)
    values = corr_triu[indices]
    return list(zip(indices[0], indices[1], values))

result = g(corr.copy())

import numpy as np
import copy

def generate_test_case(test_case_id):
    def generate_ans(data):
        corr = data
        corr_triu = np.triu(corr, k=1)
        indices = np.where(corr_triu > 0.3)
        values = corr_triu[indices]
        return list(zip(indices[0], indices[1], values))

    def define_test_input(test_case_id):
        if test_case_id == 1:
            np.random.seed(10)
            data = np.random.rand(10, 5)
            corr = np.corrcoef(data, rowvar=False)
        return corr

    test_input = define_test_input(test_case_id)
    expected_result = generate_ans(copy.deepcopy(test_input))
    return test_input, expected_result

def exec_test(result, ans):
    try:
        assert result == ans
        return 1
    except:
        return 0

exec_context = r"""
import numpy as np
corr = test_input
[insert]
"""

def test_execution(solution: str):
    code = exec_context.replace("[insert]", solution)
    for i in range(1):
        test_input, expected_result = generate_test_case(i + 1)
        test_env = {"test_input": test_input}
        exec(code, test_env)
        assert exec_test(test_env["result"], expected_result)

def g(corr):
    corr_triu = corr.where(~np.tril(np.ones(corr.shape)).astype(bool))
    corr_triu = corr_triu.stack()
    return corr_triu[corr_triu > 0.3]

result = g(corr.copy())

import pandas as pd
import numpy as np
import copy

def generate_test_case(test_case_id):
    def generate_ans(data):
        corr = data
        corr_triu = corr.where(~np.tril(np.ones(corr.shape)).astype(bool))
        corr_triu = corr_triu.stack()
        return corr_triu[corr_triu > 0.3]

    def define_test_input(test_case_id):
        if test_case_id == 1:
            np.random.seed(10)
            df = pd.DataFrame(np.random.rand(10, 5))
            corr = df.corr()
        return corr

    test_input = define_test_input(test_case_id)
    expected_result = generate_ans(copy.deepcopy(test_input))
    return test_input, expected_result

def exec_test(result, ans):
    try:
        pd.testing.assert_series_equal(result, ans, check_dtype=False)
        return 1
    except:
        return 0

exec_context = r"""
import pandas as pd
import numpy as np
corr = test_input
[insert]
"""

def test_execution(solution: str):
    code = exec_context.replace("[insert]", solution)
    for i in range(1):
        test_input, expected_result = generate_test_case(i + 1)
        test_env = {"test_input": test_input}
        exec(code, test_env)
        assert exec_test(test_env["result"], expected_result)


import numpy as np

arr = np.array([[1, 2, 3], [4, 5, 6], [7, 8, 9]])
headers = np.array(['A', 'B', 'A'])

result = ... # put solution in this variable

def g(arr, headers):
    headers[-1] = 'Test'
    return arr, headers

result = g(arr.copy(), headers.copy())

import numpy as np
import copy

def generate_test_case(test_case_id):
    def generate_ans(data, headers):
        headers[-1] = 'Test'
        return data, headers

    def define_test_input(test_case_id):
        if test_case_id == 1:
            arr = np.array([[1, 2, 3], [4, 5, 6], [7, 8, 9]])
            headers = np.array(['A', 'B', 'A'])
        return arr, headers

    test_input, headers = define_test_input(test_case_id)
    expected_result = generate_ans(copy.deepcopy(test_input), copy.deepcopy(headers))
    return (test_input, headers), expected_result

def exec_test(result, ans):
    try:
        np.testing.assert_array_equal(result[0], ans[0])
        np.testing.assert_array_equal(result[1], ans[1])
        return 1
    except:
        return 0

exec_context = r"""
import numpy as np
arr, headers = test_input
[insert]
"""

def test_execution(solution: str):
    code = exec_context.replace("[insert]", solution)
    for i in range(1):
        test_input, expected_result = generate_test_case(i + 1)
        test_env = {"test_input": test_input}
        exec(code, test_env)
        assert exec_test(test_env["result"], expected_result)

import pandas as pd

df = pd.DataFrame([[1, 2, 3], [4, 5, 6], [7, 8, 9]], columns=list('ABA'))

result = ... # put solution in this variable

def g(df):
    return df.set_axis([*df.columns[:-1], 'Test'], axis=1)

result = g(df.copy())

import pandas as pd
import numpy as np
import copy

def generate_test_case(test_case_id):
    def generate_ans(data):
        df = data
        return df.set_axis([*df.columns[:-1], "Test"], axis=1)

    def define_test_input(test_case_id):
        if test_case_id == 1:
            df = pd.DataFrame([[1, 2, 3], [4, 5, 6], [7, 8, 9]], columns=list("ABA"))
        return df

    test_input = define_test_input(test_case_id)
    expected_result = generate_ans(copy.deepcopy(test_input))
    return test_input, expected_result

def exec_test(result, ans):
    try:
        pd.testing.assert_frame_equal(result, ans, check_dtype=False)
        return 1
    except:
        return 0

exec_context = r"""
import pandas as pd
import numpy as np
df = test_input
[insert]
"""

def test_execution(solution: str):
    code = exec_context.replace("[insert]", solution)
    for i in range(1):
        test_input, expected_result = generate_test_case(i + 1)
        test_env = {"test_input": test_input}
        exec(code, test_env)
        assert exec_test(test_env["result"], expected_result)


arr = np.array([['Name', 'Name', 'Age'], ['Alice', 'Bob', '25']])

def g(arr, headers):
    headers[0] = 'Test'
    return arr, headers

result = g(arr.copy(), headers.copy())

import numpy as np
import copy

def generate_test_case(test_case_id):
    def generate_ans(data, headers):
        headers[0] = 'Test'
        return data, headers

    def define_test_input(test_case_id):
        if test_case_id == 1:
            arr = np.array([[1, 2, 3], [4, 5, 6], [7, 8, 9]])
            headers = np.array(['A', 'B', 'A'])
        return arr, headers

    test_input, headers = define_test_input(test_case_id)
    expected_result = generate_ans(copy.deepcopy(test_input), copy.deepcopy(headers))
    return (test_input, headers), expected_result

def exec_test(result, ans):
    try:
        np.testing.assert_array_equal(result[0], ans[0])
        np.testing.assert_array_equal(result[1], ans[1])
        return 1
    except:
        return 0

exec_context = r"""
import numpy as np
arr, headers = test_input
[insert]
"""

def test_execution(solution: str):
    code = exec_context.replace("[insert]", solution)
    for i in range(1):
        test_input, expected_result = generate_test_case(i + 1)
        test_env = {"test_input": test_input}
        exec(code, test_env)
        assert exec_test(test_env["result"], expected_result)

def g(df):
    return df.set_axis(['Test', *df.columns[1:]], axis=1)

result = g(df.copy())

import pandas as pd
import numpy as np
import copy

def generate_test_case(test_case_id):
    def generate_ans(data):
        df = data
        return df.set_axis(["Test", *df.columns[1:]], axis=1)

    def define_test_input(test_case_id):
        if test_case_id == 1:
            df = pd.DataFrame([[1, 2, 3], [4, 5, 6], [7, 8, 9]], columns=list("ABA"))
        return df

    test_input = define_test_input(test_case_id)
    expected_result = generate_ans(copy.deepcopy(test_input))
    return test_input, expected_result

def exec_test(result, ans):
    try:
        pd.testing.assert_frame_equal(result, ans, check_dtype=False)
        return 1
    except:
        return 0

exec_context = r"""
import pandas as pd
import numpy as np
df = test_input
[insert]
"""

def test_execution(solution: str):
    code = exec_context.replace("[insert]", solution)
    for i in range(1):
        test_input, expected_result = generate_test_case(i + 1)
        test_env = {"test_input": test_input}
        exec(code, test_env)
        assert exec_test(test_env["result"], expected_result)


import numpy as np

arr = np.array([[0, 0, 1, 1, 0],
                  [1, 1, 0, 0, 1],
                  [1, 0, 1, 1, 1]])

arr = ... # put solution in this variable

def g(arr):
    frequent = np.apply_along_axis(lambda x: np.bincount(x).argmax(), 1, arr)
    freq_count = np.apply_along_axis(lambda x: np.bincount(x).max(), 1, arr)
    return np.column_stack((arr, frequent, freq_count))

arr = g(arr.copy())

import numpy as np
import copy

def generate_test_case(test_case_id):
    def generate_ans(data):
        arr = data
        frequent = np.apply_along_axis(lambda x: np.bincount(x).argmax(), 1, arr)
        freq_count = np.apply_along_axis(lambda x: np.bincount(x).max(), 1, arr)
        return np.column_stack((arr, frequent, freq_count))

    def define_test_input(test_case_id):
        if test_case_id == 1:
            arr = np.array([
                [0, 0, 1, 1, 0],
                [1, 1, 0, 0, 1],
                [1, 0, 1, 1, 1],
            ])
        return arr

    test_input = define_test_input(test_case_id)
    expected_result = generate_ans(copy.deepcopy(test_input))
    return test_input, expected_result

def exec_test(result, ans):
    try:
        np.testing.assert_array_equal(result, ans)
        return 1
    except:
        return 0

exec_context = r"""
import numpy as np
arr = test_input
[insert]
result = arr
"""

def test_execution(solution: str):
    code = exec_context.replace("[insert]", solution)
    for i in range(1):
        test_input, expected_result = generate_test_case(i + 1)
        test_env = {"test_input": test_input}
        exec(code, test_env)
        assert exec_test(test_env["result"], expected_result)

import pandas as pd

df = pd.DataFrame({'bit1': [0, 1, 1],
                   'bit2': [0, 1, 0],
                   'bit3': [1, 0, 1],
                   'bit4': [1, 0, 1],
                   'bit5': [0, 1, 1]})

df = ... # put solution in this variable

def g(df):
    df['frequent'] = df.mode(axis=1)
    for i in df.index:
        df.loc[i, 'freq_count'] = (df.iloc[i]==df.loc[i, 'frequent']).sum() - 1
    return df

df = g(df.copy())

import pandas as pd
import numpy as np
import copy

def generate_test_case(test_case_id):
    def generate_ans(data):
        df = data
        df["frequent"] = df.mode(axis=1)
        for i in df.index:
            df.loc[i, "freq_count"] = (df.iloc[i] == df.loc[i, "frequent"]).sum() - 1
        return df

    def define_test_input(test_case_id):
        if test_case_id == 1:
            df = pd.DataFrame(
                {
                    "bit1": [0, 1, 1],
                    "bit2": [0, 1, 0],
                    "bit3": [1, 0, 1],
                    "bit4": [1, 0, 1],
                    "bit5": [0, 1, 1],
                }
            )
        return df

    test_input = define_test_input(test_case_id)
    expected_result = generate_ans(copy.deepcopy(test_input))
    return test_input, expected_result

def exec_test(result, ans):
    try:
        pd.testing.assert_frame_equal(result, ans)
        return 1
    except:
        return 0

exec_context = r"""
import pandas as pd
import numpy as np
df = test_input
[insert]
result = df
"""

def test_execution(solution: str):
    code = exec_context.replace("[insert]", solution)
    for i in range(1):
        test_input, expected_result = generate_test_case(i + 1)
        test_env = {"test_input": test_input}
        exec(code, test_env)
        assert exec_test(test_env["result"], expected_result)


import numpy as np

arr = np.array([[0, 0, 3, 3, 0],
                  [2, 2, 0, 0, 2],
                  [4, 0, 4, 4, 4]])

arr = ... # put solution in this variable

import numpy as np
import copy

def generate_test_case(test_case_id):
    def generate_ans(data):
        arr = data
        frequent = np.apply_along_axis(lambda x: np.bincount(x).argmax(), 1, arr)
        freq_count = np.apply_along_axis(lambda x: np.bincount(x).max(), 1, arr)
        return np.column_stack((arr, frequent, freq_count))

    def define_test_input(test_case_id):
        if test_case_id == 1:
            arr = np.array([
                [0, 0, 3, 3, 0],
                [2, 2, 0, 0, 2],
                [4, 0, 4, 4, 4],
            ])
        return arr

    test_input = define_test_input(test_case_id)
    expected_result = generate_ans(copy.deepcopy(test_input))
    return test_input, expected_result

def exec_test(result, ans):
    try:
        np.testing.assert_array_equal(result, ans)
        return 1
    except:
        return 0

exec_context = r"""
import numpy as np
arr = test_input
[insert]
result = arr
"""

def test_execution(solution: str):
    code = exec_context.replace("[insert]", solution)
    for i in range(1):
        test_input, expected_result = generate_test_case(i + 1)
        test_env = {"test_input": test_input}
        exec(code, test_env)
        assert exec_test(test_env["result"], expected_result)

import pandas as pd

df = pd.DataFrame({'bit1': [0, 2, 4],
                   'bit2': [0, 2, 0],
                   'bit3': [3, 0, 4],
                   'bit4': [3, 0, 4],
                   'bit5': [0, 2, 4]})

df = ... # put solution in this variable

import pandas as pd
import numpy as np
import copy

def generate_test_case(test_case_id):
    def generate_ans(data):
        df = data
        df["frequent"] = df.mode(axis=1)
        for i in df.index:
            df.loc[i, "freq_count"] = (df.iloc[i] == df.loc[i, "frequent"]).sum() - 1
        return df

    def define_test_input(test_case_id):
        if test_case_id == 1:
            df = pd.DataFrame(
                {
                    "bit1": [0, 2, 4],
                    "bit2": [0, 2, 0],
                    "bit3": [3, 0, 4],
                    "bit4": [3, 0, 4],
                    "bit5": [0, 2, 4],
                }
            )
        return df

    test_input = define_test_input(test_case_id)
    expected_result = generate_ans(copy.deepcopy(test_input))
    return test_input, expected_result

def exec_test(result, ans):
    try:
        pd.testing.assert_frame_equal(result, ans)
        return 1
    except:
        return 0

exec_context = r"""
import pandas as pd
import numpy as np
df = test_input
[insert]
result = df
"""

def test_execution(solution: str):
    code = exec_context.replace("[insert]", solution)
    for i in range(1):
        test_input, expected_result = generate_test_case(i + 1)
        test_env = {"test_input": test_input}
        exec(code, test_env)
        assert exec_test(test_env["result"], expected_result)


import numpy as np

arr = np.array([[0, 2, 4, 3, 0, 3],
                  [0, 2, 0, 0, 2, 0],
                  [3, 0, 4, 4, 4, 5]])

result = ... # put solution in this variable

def find_frequent(arr):
    frequent = []
    freq_count = []
    for row in arr:
        values, counts = np.unique(row, return_counts=True)
        max_count = np.max(counts)
        frequent_values = values[counts == max_count]
        frequent.append(sorted(frequent_values.tolist()))
        freq_count.append(max_count)
    return np.array(frequent, dtype=object), np.array(freq_count)

frequent, freq_count = find_frequent(arr)

import numpy as np
import copy

def generate_test_case(test_case_id):
    def generate_ans(data):
        arr = data
        frequent = []
        freq_count = []
        for row in arr:
            values, counts = np.unique(row, return_counts=True)
            max_count = np.max(counts)
            frequent_values = values[counts == max_count]
            frequent.append(sorted(frequent_values.tolist()))
            freq_count.append(max_count)
        return np.array(frequent, dtype=object), np.array(freq_count)

    def define_test_input(test_case_id):
        if test_case_id == 1:
            arr = np.array([
                [0, 2, 4, 3, 0, 3],
                [0, 2, 0, 0, 2, 0],
                [3, 0, 4, 4, 4, 5],
            ])
        return arr

    test_input = define_test_input(test_case_id)
    expected_result = generate_ans(copy.deepcopy(test_input))
    return test_input, expected_result

def exec_test(result, ans):
    try:
        assert np.array_equal(result[0], ans[0]) and np.array_equal(result[1], ans[1])
        return 1
    except:
        return 0

exec_context = r"""
import numpy as np
arr = test_input
[insert]
result = (frequent, freq_count)
"""

def test_execution(solution: str):
    code = exec_context.replace("[insert]", solution)
    for i in range(1):
        test_input, expected_result = generate_test_case(i + 1)
        test_env = {"test_input": test_input}
        exec(code, test_env)
        assert exec_test(test_env["result"], expected_result)

def g(df):
    cols = list(df)
    Mode = df.mode(axis=1)
    df['frequent'] = df['bit1'].astype(object)
    for i in df.index:
        df.at[i, 'frequent'] = []
    for i in df.index:
        for col in list(Mode):
            if pd.isna(Mode.loc[i, col])==False:
                df.at[i, 'frequent'].append(Mode.loc[i, col])
        df.at[i, 'frequent'] = sorted(df.at[i, 'frequent'])
        df.loc[i, 'freq_count'] = (df[cols].iloc[i]==df.loc[i, 'frequent'][0]).sum()
    return df

df = g(df.copy())

import pandas as pd
import numpy as np
import copy

def generate_test_case(test_case_id):
    def generate_ans(data):
        df = data
        cols = list(df)
        Mode = df.mode(axis=1)
        df["frequent"] = df["bit1"].astype(object)
        for i in df.index:
            df.at[i, "frequent"] = []
        for i in df.index:
            for col in list(Mode):
                if pd.isna(Mode.loc[i, col]) == False:
                    df.at[i, "frequent"].append(Mode.loc[i, col])
            df.at[i, "frequent"] = sorted(df.at[i, "frequent"])
            df.loc[i, "freq_count"] = (
                df[cols].iloc[i] == df.loc[i, "frequent"][0]
            ).sum()
        return df

    def define_test_input(test_case_id):
        if test_case_id == 1:
            df = pd.DataFrame(
                {
                    "bit1": [0, 2, 4],
                    "bit2": [0, 2, 0],
                    "bit3": [3, 0, 4],
                    "bit4": [3, 0, 4],
                    "bit5": [0, 2, 4],
                    "bit6": [3, 0, 5],
                }
            )
        return df

    test_input = define_test_input(test_case_id)
    expected_result = generate_ans(copy.deepcopy(test_input))
    return test_input, expected_result

def exec_test(result, ans):
    try:
        pd.testing.assert_frame_equal(result, ans)
        return 1
    except:
        return 0

exec_context = r"""
import pandas as pd
import numpy as np
df = test_input
[insert]
for i in df.index:
    df.at[i, 'frequent'] = sorted(df.at[i, 'frequent'])
result = df
"""

def test_execution(solution: str):
    code = exec_context.replace("[insert]", solution)
    for i in range(1):
        test_input, expected_result = generate_test_case(i + 1)
        test_env = {"test_input": test_input}
        exec(code, test_env)
        assert exec_test(test_env["result"], expected_result)


def g(arr):
    arr[arr == 'NULL'] = 0
    arr = arr.astype(float)
    unique_keys, indices = np.unique(arr[:, [1, 3]], axis=0, return_inverse=True)
    sums = np.zeros((unique_keys.shape[0], 2))
    counts = np.zeros((unique_keys.shape[0], 2))
    for i, idx in enumerate(indices):
        sums[idx] += arr[i, [0, 2]]
        counts[idx] += 1
    means = sums / counts
    return np.hstack((unique_keys, means))

result = g(arr.copy())

import numpy as np
import copy

def generate_test_case(test_case_id):
    def generate_ans(data):
        arr = data
        arr[arr == 'NULL'] = 0
        arr = arr.astype(float)
        unique_keys, indices = np.unique(arr[:, [1, 3]], axis=0, return_inverse=True)
        sums = np.zeros((unique_keys.shape[0], 2))
        counts = np.zeros((unique_keys.shape[0], 2))
        for i, idx in enumerate(indices):
            sums[idx] += arr[i, [0, 2]]
            counts[idx] += 1
        means = sums / counts
        return np.hstack((unique_keys, means))

    def define_test_input(test_case_id):
        if test_case_id == 1:
            arr = np.array([
                [8.0, 1, 'NULL', 1],
                [5.0, 1, 'NULL', 1],
                [3.0, 1, 'NULL', 1],
                [4.0, 1, 1, 2],
                [7.0, 1, 3, 2],
                [9.0, 1, 4, 3],
                [5.0, 1, 2, 3],
                [7.0, 1, 3, 1]
            ])
        if test_case_id == 2:
            arr = np.array([
                [18.0, 1, 'NULL', 1],
                [5.0, 1, 'NULL', 1],
                [3.0, 1, 'NULL', 1],
                [4.0, 1, 11, 2],
                [17.0, 1, 3, 2],
                [9.0, 1, 4, 3],
                [5.0, 1, 2, 3],
                [7.0, 1, 3, 1]
            ])
        return arr

    test_input = define_test_input(test_case_id)
    expected_result = generate_ans(copy.deepcopy(test_input))
    return test_input, expected_result

def exec_test(result, ans):
    try:
        np.testing.assert_array_almost_equal(result, ans)
        return 1
    except:
        return 0

exec_context = r"""
import numpy as np
arr = test_input
[insert]
"""

def test_execution(solution: str):
    code = exec_context.replace("[insert]", solution)
    for i in range(2):
        test_input, expected_result = generate_test_case(i + 1)
        test_env = {"test_input": test_input}
        exec(code, test_env)
        assert exec_test(test_env["result"], expected_result)

def g(df):
    df['bar'] = df['bar'].replace("NULL", 0)
    res = df.groupby(["id1", "id2"])[["foo", "bar"]].mean()
    return res

result = g(df.copy())

import pandas as pd
import numpy as np
import copy

def generate_test_case(test_case_id):
    def generate_ans(data):
        df = data
        df["bar"] = df["bar"].replace("NULL", 0)
        res = df.groupby(["id1", "id2"])[["foo", "bar"]].mean()
        return res

    def define_test_input(test_case_id):
        if test_case_id == 1:
            df = pd.DataFrame(
                {
                    "foo": [8, 5, 3, 4, 7, 9, 5, 7],
                    "id1": [1, 1, 1, 1, 1, 1, 1, 1],
                    "bar": ["NULL", "NULL", "NULL", 1, 3, 4, 2, 3],
                    "id2": [1, 1, 1, 2, 2, 3, 3, 1],
                }
            )
        if test_case_id == 2:
            df = pd.DataFrame(
                {
                    "foo": [18, 5, 3, 4, 17, 9, 5, 7],
                    "id1": [1, 1, 1, 1, 1, 1, 1, 1],
                    "bar": ["NULL", "NULL", "NULL", 11, 3, 4, 2, 3],
                    "id2": [1, 1, 1, 2, 2, 3, 3, 1],
                }
            )
        return df

    test_input = define_test_input(test_case_id)
    expected_result = generate_ans(copy.deepcopy(test_input))
    return test_input, expected_result

def exec_test(result, ans):
    try:
        pd.testing.assert_frame_equal(result, ans, check_dtype=False)
        return 1
    except:
        return 0

exec_context = r"""
import pandas as pd
import numpy as np
df = test_input
[insert]
"""

def test_execution(solution: str):
    code = exec_context.replace("[insert]", solution)
    for i in range(2):
        test_input, expected_result = generate_test_case(i + 1)
        test_env = {"test_input": test_input}
        exec(code, test_env)
        assert exec_test(test_env["result"], expected_result)


def g(df_a, df_b):
    result = []
    for row_a in df_a:
        for row_b in df_b:
            if row_a[0] == row_b[0]:
                result.append([row_a[0], row_a[1], row_b[1]])
    return np.array(result)

result = g(df_a.copy(), df_b.copy())

import numpy as np
import copy

def generate_test_case(test_case_id):
    def generate_ans(data):
        data = data
        df_a, df_b = data
        result = []
        for row_a in df_a:
            for row_b in df_b:
                if row_a[0] == row_b[0]:
                    result.append([row_a[0], row_a[1], row_b[1]])
        return np.array(result)

    def define_test_input(test_case_id):
        if test_case_id == 1:
            df_a = np.array([[1001.01, 100], [1002.02, 50], [1003.03, 200]])
            df_b = np.array([[1001.01, 'alice', 7], [1002.02, 'bob', 8], [1003.03, '777', 9]])
        if test_case_id == 2:
            df_a = np.array([[1001.01, 100], [1002.02, 50], [1003.03, 200]])
            df_b = np.array([[1001.01, '666', 7], [1002.02, 'bob', 8], [1003.03, 'alice', 9]])
        return df_a, df_b

    test_input = define_test_input(test_case_id)
    expected_result = generate_ans(copy.deepcopy(test_input))
    return test_input, expected_result

def exec_test(result, ans):
    try:
        np.testing.assert_array_equal(result, ans)
        return 1
    except:
        return 0

exec_context = r"""
import numpy as np
df_a, df_b = test_input
[insert]
"""

def test_execution(solution: str):
    code = exec_context.replace("[insert]", solution)
    for i in range(2):
        test_input, expected_result = generate_test_case(i + 1)
        test_env = {"test_input": test_input}
        exec(code, test_env)
        assert exec_test(test_env["result"], expected_result)

def g(df_a, df_b):
    return df_a[['EntityNum', 'foo']].merge(df_b[['EntityNum', 'a_col']], on='EntityNum', how='left')

result = g(df_a.copy(), df_b.copy())

import pandas as pd
import numpy as np
import copy

def generate_test_case(test_case_id):
    def generate_ans(data):
        data = data
        df_a, df_b = data
        return df_a[["EntityNum", "foo"]].merge(
            df_b[["EntityNum", "a_col"]], on="EntityNum", how="left"
        )

    def define_test_input(test_case_id):
        if test_case_id == 1:
            df_a = pd.DataFrame(
                {"EntityNum": [1001.01, 1002.02, 1003.03], "foo": [100, 50, 200]}
            )
            df_b = pd.DataFrame(
                {
                    "EntityNum": [1001.01, 1002.02, 1003.03],
                    "a_col": ["alice", "bob", "777"],
                    "b_col": [7, 8, 9],
                }
            )
        if test_case_id == 2:
            df_a = pd.DataFrame(
                {"EntityNum": [1001.01, 1002.02, 1003.03], "foo": [100, 50, 200]}
            )
            df_b = pd.DataFrame(
                {
                    "EntityNum": [1001.01, 1002.02, 1003.03],
                    "a_col": ["666", "bob", "alice"],
                    "b_col": [7, 8, 9],
                }
            )
        return df_a, df_b

    test_input = define_test_input(test_case_id)
    expected_result = generate_ans(copy.deepcopy(test_input))
    return test_input, expected_result

def exec_test(result, ans):
    try:
        pd.testing.assert_frame_equal(result, ans, check_dtype=False)
        return 1
    except:
        return 0

exec_context = r"""
import pandas as pd
import numpy as np
df_a, df_b = test_input
[insert]
"""

def test_execution(solution: str):
    code = exec_context.replace("[insert]", solution)
    for i in range(2):
        test_input, expected_result = generate_test_case(i + 1)
        test_env = {"test_input": test_input}
        exec(code, test_env)
        assert exec_test(test_env["result"], expected_result)