Prompt:
Problem:
I need to do some analysis on a large dataset from a hydrolgeology field work. I am using NumPy. I want to know how I can:
1.	divide e.g. the row-th row of my array by a number (e.g. 5.2). And then
2.	calculate the multiplication of the numbers in that row.
As I mentioned I only want to work on a specific row and not the whole array. The result should be that of multiplication
A:
<code>
import numpy as np
a = np.random.rand(8, 5)
row = 2
divide_number = 5.2
</code>
result = ... # put solution in this variable
BEGIN SOLUTION
<code>

a[row-1, :] /= divide_number
result = np.multiply.reduce(a[row-1, :])

import numpy as np
import copy

def generate_test_case(test_case_id):
    def define_test_input(test_case_id):
        if test_case_id == 1:
            np.random.seed(42)
            a = np.random.rand(8, 5)
            row = 2
            const = 5.2
        elif test_case_id == 2:
            np.random.seed(42)
            a = np.random.rand(np.random.randint(5, 10), np.random.randint(6, 10))
            row = 4
            const = np.random.rand() + 1
        return a, row, const

    def generate_ans(data):
        _a = data
        a, row, divide_number = _a
        a[row - 1, :] /= divide_number
        result = np.multiply.reduce(a[row - 1, :])
        return result

    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):
    np.testing.assert_allclose(result, ans)
    return 1

exec_context = r"""
import numpy as np
a, row, divide_number = 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)


Prompt:
Problem:
How to get one maximal set of linearly independent vectors of a given matrix `a`?
For example, [[0 1 0 0], [0 0 1 0], [1 0 0 1]] in [[0 1 0 0], [0 0 1 0], [0 1 1 0], [1 0 0 1]]
A:
<code>
import numpy as np
a = np.array([[0,1,0,0], [0,0,1,0], [0,1,1,0], [1,0,0,1]])
</code>
result = ... # put solution in this variable
BEGIN SOLUTION
<code>

def LI_vecs(M):
    dim = M.shape[0]
    LI=[M[0]]
    for i in range(dim):
        tmp=[]
        for r in LI:
            tmp.append(r)
        tmp.append(M[i])                #set tmp=LI+[M[i]]
        if np.linalg.matrix_rank(tmp)>len(LI):    #test if M[i] is linearly independent from all (row) vectors in LI
            LI.append(M[i])             #note that matrix_rank does not need to take in a square matrix
    return LI                           #return set of linearly independent (row) vectors
result = LI_vecs(a)

import numpy as np
import pandas as pd
import copy

def generate_test_case(test_case_id):
    def define_test_input(test_case_id):
        if test_case_id == 1:
            a = np.array([[0, 1, 0, 0], [0, 0, 1, 0], [0, 1, 1, 0], [1, 0, 0, 1]])
        elif test_case_id == 2:
            a = np.array(
                [
                    [0, 1, 0, 0, 1, 1],
                    [0, 0, 1, 0, 0, 1],
                    [0, 1, 1, 0, 0, 0],
                    [1, 0, 1, 0, 0, 1],
                    [1, 1, 0, 1, 0, 0],
                ]
            )
        elif test_case_id == 3:
            a = np.array(
                [
                    [0, 1, 0, 0, 1, 1],
                    [0, 1, 0, 0, 1, 1],
                    [0, 1, 0, 0, 1, 1],
                    [1, 0, 1, 0, 0, 1],
                    [1, 1, 0, 1, 0, 0],
                ]
            )
        return a

    def generate_ans(data):
        _a = data
        a = _a
        return a

    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):
    result = np.array(result)
    if result.shape[1] == ans.shape[1]:
        assert np.linalg.matrix_rank(ans) == np.linalg.matrix_rank(
            result
        ) and np.linalg.matrix_rank(result) == len(result)
        assert len(np.unique(result, axis=0)) == len(result)
        for arr in result:
            assert np.any(np.all(ans == arr, axis=1))
    else:
        assert (
            np.linalg.matrix_rank(ans) == np.linalg.matrix_rank(result)
            and np.linalg.matrix_rank(result) == result.shape[1]
        )
        assert np.unique(result, axis=1).shape[1] == result.shape[1]
        for i in range(result.shape[1]):
            assert np.any(np.all(ans == result[:, i].reshape(-1, 1), axis=0))
    return 1

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

def test_execution(solution: str):
    code = exec_context.replace("[insert]", solution)
    for i in range(3):
        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)


Prompt:
Problem:
How do i get the length of the row in a 2D array?
example, i have a nD array called a. when i print a.shape, it returns (1,21). I want to do a for loop, in the range of the row size (21) of the array a. How do i get the value of row size as result?
A:
<code>
import numpy as np
a = np.random.rand(np.random.randint(5, 10), np.random.randint(6, 10))
</code>
result = ... # put solution in this variable
BEGIN SOLUTION
<code>

import numpy as np
import pandas as pd
import copy

def generate_test_case(test_case_id):
    def define_test_input(test_case_id):
        if test_case_id == 1:
            np.random.seed(42)
            a = np.random.rand(np.random.randint(5, 10), np.random.randint(6, 10))
        return a

    def generate_ans(data):
        _a = data
        a = _a
        result = a.shape[1]
        return result

    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):
    np.testing.assert_array_equal(result, ans)
    return 1

exec_context = r"""
import numpy as np
a = 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)


Prompt:
Problem:
I have data of sample 1 and sample 2 (`a` and `b`) – size is different for sample 1 and sample 2. I want to do a weighted (take n into account) two-tailed t-test.
I tried using the scipy.stat module by creating my numbers with np.random.normal, since it only takes data and not stat values like mean and std dev (is there any way to use these values directly). But it didn't work since the data arrays has to be of equal size.
Any help on how to get the p-value would be highly appreciated.
A:
<code>
import numpy as np
import scipy.stats
a = np.random.randn(40)
b = 4*np.random.randn(50)
</code>
p_value = ... # put solution in this variable
BEGIN SOLUTION
<code>

_, p_value = scipy.stats.ttest_ind(a, b,  equal_var = False)

import numpy as np
import copy
import scipy.stats

def generate_test_case(test_case_id):
    def define_test_input(test_case_id):
        if test_case_id == 1:
            np.random.seed(42)
            a = np.random.randn(40)
            b = 4 * np.random.randn(50)
        return a, b

    def generate_ans(data):
        _a = data
        a, b = _a
        _, p_value = scipy.stats.ttest_ind(a, b, equal_var=False)
        return p_value

    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):
    assert abs(ans - result) <= 1e-5
    return 1

exec_context = r"""
import numpy as np
import scipy.stats
a, b = test_input
[insert]
result = p_value
"""

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)


Prompt:
Problem:
I have data of sample 1 and sample 2 (`a` and `b`) – size is different for sample 1 and sample 2. I want to do a weighted (take n into account) two-tailed t-test.
I tried using the scipy.stat module by creating my numbers with np.random.normal, since it only takes data and not stat values like mean and std dev (is there any way to use these values directly). But it didn't work since the data arrays has to be of equal size.
For some reason, nans might be in original data, and we want to omit them.
Any help on how to get the p-value would be highly appreciated.
A:
<code>
import numpy as np
import scipy.stats
a = np.random.randn(40)
b = 4*np.random.randn(50)
</code>
p_value = ... # put solution in this variable
BEGIN SOLUTION
<code>

_, p_value = scipy.stats.ttest_ind(a, b,  equal_var = False, nan_policy = 'omit')

import numpy as np
import copy
import scipy.stats

def generate_test_case(test_case_id):
    def define_test_input(test_case_id):
        if test_case_id == 1:
            np.random.seed(42)
            a = np.random.randn(40)
            b = 4 * np.random.randn(50)
        elif test_case_id == 2:
            np.random.seed(43)
            a = np.random.randn(40)
            b = 4 * np.random.randn(50)
            a[10] = np.nan
            b[20] = np.nan
        return a, b

    def generate_ans(data):
        _a = data
        a, b = _a
        _, p_value = scipy.stats.ttest_ind(a, b, equal_var=False, nan_policy="omit")
        return p_value

    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):
    assert abs(ans - result) <= 1e-5
    return 1

exec_context = r"""
import numpy as np
import scipy.stats
a, b = test_input
[insert]
result = p_value
"""

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)


Prompt:
Problem:
I have only the summary statistics of sample 1 and sample 2, namely mean, variance, nobs(number of observations). I want to do a weighted (take n into account) two-tailed t-test.
Any help on how to get the p-value would be highly appreciated.
A:
<code>
import numpy as np
import scipy.stats
amean = -0.0896
avar = 0.954
anobs = 40
bmean = 0.719
bvar = 11.87
bnobs = 50
</code>
p_value = ... # put solution in this variable
BEGIN SOLUTION
<code>

_, p_value = scipy.stats.ttest_ind_from_stats(amean, np.sqrt(avar), anobs, bmean, np.sqrt(bvar), bnobs, equal_var=False)

import numpy as np
import copy
import scipy.stats

def generate_test_case(test_case_id):
    def define_test_input(test_case_id):
        if test_case_id == 1:
            amean = -0.0896
            avar = 0.954
            anobs = 40
            bmean = 0.719
            bvar = 11.87
            bnobs = 50
        return amean, avar, anobs, bmean, bvar, bnobs

    def generate_ans(data):
        _a = data
        amean, avar, anobs, bmean, bvar, bnobs = _a
        _, p_value = scipy.stats.ttest_ind_from_stats(
            amean, np.sqrt(avar), anobs, bmean, np.sqrt(bvar), bnobs, equal_var=False
        )
        return p_value

    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):
    assert abs(ans - result) <= 1e-5
    return 1

exec_context = r"""
import numpy as np
import scipy.stats
amean, avar, anobs, bmean, bvar, bnobs = test_input
[insert]
result = p_value
"""

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)


Prompt:
Problem:
Say I have these 2D arrays A and B.
How can I remove elements from A that are in B. (Complement in set theory: A-B)
Example:
A=np.asarray([[1,1,1], [1,1,2], [1,1,3], [1,1,4]])
B=np.asarray([[0,0,0], [1,0,2], [1,0,3], [1,0,4], [1,1,0], [1,1,1], [1,1,4]])
#in original order
#output = [[1,1,2], [1,1,3]]

A:
<code>
import numpy as np
A=np.asarray([[1,1,1], [1,1,2], [1,1,3], [1,1,4]])
B=np.asarray([[0,0,0], [1,0,2], [1,0,3], [1,0,4], [1,1,0], [1,1,1], [1,1,4]])
</code>
output = ... # put solution in this variable
BEGIN SOLUTION
<code>

dims = np.maximum(B.max(0),A.max(0))+1
output = A[~np.in1d(np.ravel_multi_index(A.T,dims),np.ravel_multi_index(B.T,dims))]

import numpy as np
import copy

def generate_test_case(test_case_id):
    def define_test_input(test_case_id):
        if test_case_id == 1:
            A = np.asarray([[1, 1, 1], [1, 1, 2], [1, 1, 3], [1, 1, 4]])
            B = np.asarray(
                [
                    [0, 0, 0],
                    [1, 0, 2],
                    [1, 0, 3],
                    [1, 0, 4],
                    [1, 1, 0],
                    [1, 1, 1],
                    [1, 1, 4],
                ]
            )
        elif test_case_id == 2:
            np.random.seed(42)
            A = np.random.randint(0, 2, (10, 3))
            B = np.random.randint(0, 2, (20, 3))
        elif test_case_id == 3:
            A = np.asarray([[1, 1, 1], [1, 1, 4]])
            B = np.asarray(
                [
                    [0, 0, 0],
                    [1, 0, 2],
                    [1, 0, 3],
                    [1, 0, 4],
                    [1, 1, 0],
                    [1, 1, 1],
                    [1, 1, 4],
                ]
            )
        return A, B

    def generate_ans(data):
        _a = data
        A, B = _a
        dims = np.maximum(B.max(0), A.max(0)) + 1
        output = A[
            ~np.in1d(np.ravel_multi_index(A.T, dims), np.ravel_multi_index(B.T, dims))
        ]
        return output

    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):
    if ans.shape[0]:
        np.testing.assert_array_equal(result, ans)
    else:
        result = result.reshape(0)
        ans = ans.reshape(0)
        np.testing.assert_array_equal(result, ans)
    return 1

exec_context = r"""
import numpy as np
A, B = test_input
[insert]
result = output
"""

def test_execution(solution: str):
    code = exec_context.replace("[insert]", solution)
    for i in range(3):
        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)


Prompt:
Problem:
Say I have these 2D arrays A and B.
How can I get elements from A that are not in B, and those from B that are not in A? (Symmetric difference in set theory: A△B)
Example:
A=np.asarray([[1,1,1], [1,1,2], [1,1,3], [1,1,4]])
B=np.asarray([[0,0,0], [1,0,2], [1,0,3], [1,0,4], [1,1,0], [1,1,1], [1,1,4]])
#elements in A first, elements in B then. in original order.
#output = array([[1,1,2], [1,1,3], [0,0,0], [1,0,2], [1,0,3], [1,0,4], [1,1,0]])

A:
<code>
import numpy as np
A=np.asarray([[1,1,1], [1,1,2], [1,1,3], [1,1,4]])
B=np.asarray([[0,0,0], [1,0,2], [1,0,3], [1,0,4], [1,1,0], [1,1,1], [1,1,4]])
</code>
output = ... # put solution in this variable
BEGIN SOLUTION
<code>

dims = np.maximum(B.max(0),A.max(0))+1
result = A[~np.in1d(np.ravel_multi_index(A.T,dims),np.ravel_multi_index(B.T,dims))]
output = np.append(result, B[~np.in1d(np.ravel_multi_index(B.T,dims),np.ravel_multi_index(A.T,dims))], axis = 0)

import numpy as np
import pandas as pd
import copy

def generate_test_case(test_case_id):
    def define_test_input(test_case_id):
        if test_case_id == 1:
            A = np.asarray([[1, 1, 1], [1, 1, 2], [1, 1, 3], [1, 1, 4]])
            B = np.asarray(
                [
                    [0, 0, 0],
                    [1, 0, 2],
                    [1, 0, 3],
                    [1, 0, 4],
                    [1, 1, 0],
                    [1, 1, 1],
                    [1, 1, 4],
                ]
            )
        elif test_case_id == 2:
            np.random.seed(42)
            A = np.random.randint(0, 2, (10, 3))
            B = np.random.randint(0, 2, (20, 3))
        elif test_case_id == 3:
            A = np.asarray([[1, 1, 1], [1, 1, 4]])
            B = np.asarray(
                [
                    [0, 0, 0],
                    [1, 0, 2],
                    [1, 0, 3],
                    [1, 0, 4],
                    [1, 1, 0],
                    [1, 1, 1],
                    [1, 1, 4],
                ]
            )
        return A, B

    def generate_ans(data):
        _a = data
        A, B = _a
        dims = np.maximum(B.max(0), A.max(0)) + 1
        result = A[
            ~np.in1d(np.ravel_multi_index(A.T, dims), np.ravel_multi_index(B.T, dims))
        ]
        output = np.append(
            result,
            B[
                ~np.in1d(
                    np.ravel_multi_index(B.T, dims), np.ravel_multi_index(A.T, dims)
                )
            ],
            axis=0,
        )
        return output

    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):
    if ans.shape[0]:
        np.testing.assert_array_equal(result, ans)
    else:
        result = result.reshape(0)
        ans = ans.reshape(0)
        np.testing.assert_array_equal(result, ans)
    return 1

exec_context = r"""
import numpy as np
A, B = test_input
[insert]
result = output
"""

def test_execution(solution: str):
    code = exec_context.replace("[insert]", solution)
    for i in range(3):
        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)


Prompt:
Problem:
Similar to this answer, I have a pair of 3D numpy arrays, a and b, and I want to sort the entries of b by the values of a. Unlike this answer, I want to sort only along one axis of the arrays.
My naive reading of the numpy.argsort() documentation:
Returns
-------
index_array : ndarray, int
    Array of indices that sort `a` along the specified axis.
    In other words, ``a[index_array]`` yields a sorted `a`.
led me to believe that I could do my sort with the following code:
import numpy
print a
"""
[[[ 1.  1.  1.]
  [ 1.  1.  1.]
  [ 1.  1.  1.]]
 [[ 3.  3.  3.]
  [ 3.  2.  3.]
  [ 3.  3.  3.]]
 [[ 2.  2.  2.]
  [ 2.  3.  2.]
  [ 2.  2.  2.]]]
"""
b = numpy.arange(3*3*3).reshape((3, 3, 3))
print "b"
print b
"""
[[[ 0  1  2]
  [ 3  4  5]
  [ 6  7  8]]
 [[ 9 10 11]
  [12 13 14]
  [15 16 17]]
 [[18 19 20]
  [21 22 23]
  [24 25 26]]]
##This isnt' working how I'd like
sort_indices = numpy.argsort(a, axis=0)
c = b[sort_indices]
"""
Desired output:
[[[ 0  1  2]
  [ 3  4  5]
  [ 6  7  8]]
 [[18 19 20]
  [21 13 23]
  [24 25 26]]
 [[ 9 10 11]
  [12 22 14]
  [15 16 17]]]
"""
print "Desired shape of b[sort_indices]: (3, 3, 3)."
print "Actual shape of b[sort_indices]:"
print c.shape
"""
(3, 3, 3, 3, 3)
"""
What's the right way to do this?
A:
<code>
import numpy as np
a = np.random.rand(3, 3, 3)
b = np.arange(3*3*3).reshape((3, 3, 3))
</code>
c = ... # put solution in this variable
BEGIN SOLUTION
<code>

import numpy as np
import copy

def generate_test_case(test_case_id):
    def define_test_input(test_case_id):
        if test_case_id == 1:
            np.random.seed(42)
            a = np.random.rand(3, 3, 3)
            b = np.arange(3 * 3 * 3).reshape((3, 3, 3))
        return a, b

    def generate_ans(data):
        _a = data
        a, b = _a
        sort_indices = np.argsort(a, axis=0)
        static_indices = np.indices(a.shape)
        c = b[sort_indices, static_indices[1], static_indices[2]]
        return c

    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):
    np.testing.assert_array_equal(result, ans)
    return 1

exec_context = r"""
import numpy as np
a, b = test_input
[insert]
result = c
"""

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)


Prompt:
Problem:
Similar to this answer, I have a pair of 3D numpy arrays, a and b, and I want to sort the entries of b by the values of a. Unlike this answer, I want to sort only along one axis of the arrays.
My naive reading of the numpy.argsort() documentation:
Returns
-------
index_array : ndarray, int
    Array of indices that sort `a` along the specified axis.
    In other words, ``a[index_array]`` yields a sorted `a`.
led me to believe that I could do my sort with the following code:
import numpy
print a
"""
[[[ 1.  1.  1.]
  [ 1.  1.  1.]
  [ 1.  1.  1.]]
 [[ 3.  3.  3.]
  [ 3.  3.  3.]
  [ 3.  3.  3.]]
 [[ 2.  2.  2.]
  [ 2.  2.  2.]
  [ 2.  2.  2.]]]
"""
b = numpy.arange(3*3*3).reshape((3, 3, 3))
print "b"
print b
"""
[[[ 0  1  2]
  [ 3  4  5]
  [ 6  7  8]]
 [[ 9 10 11]
  [12 13 14]
  [15 16 17]]
 [[18 19 20]
  [21 22 23]
  [24 25 26]]]
##This isnt' working how I'd like
sort_indices = numpy.argsort(a, axis=0)
c = b[sort_indices]
"""
Desired output:
[[[ 0  1  2]
  [ 3  4  5]
  [ 6  7  8]]
 [[18 19 20]
  [21 22 23]
  [24 25 26]]
 [[ 9 10 11]
  [12 13 14]
  [15 16 17]]]
"""
print "Desired shape of b[sort_indices]: (3, 3, 3)."
print "Actual shape of b[sort_indices]:"
print c.shape
"""
(3, 3, 3, 3, 3)
"""
What's the right way to do this?
A:
<code>
import numpy as np
a = np.random.rand(3, 3, 3)
b = np.arange(3*3*3).reshape((3, 3, 3))
</code>
c = ... # put solution in this variable
BEGIN SOLUTION
<code>

import numpy as np
import copy

def generate_test_case(test_case_id):
    def define_test_input(test_case_id):
        if test_case_id == 1:
            np.random.seed(42)
            a = np.random.rand(3, 3, 3)
            b = np.arange(3 * 3 * 3).reshape((3, 3, 3))
        return a, b

    def generate_ans(data):
        _a = data
        a, b = _a
        sort_indices = np.argsort(a, axis=0)
        static_indices = np.indices(a.shape)
        c = b[sort_indices, static_indices[1], static_indices[2]]
        return c

    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):
    np.testing.assert_array_equal(result, ans)
    return 1

exec_context = r"""
import numpy as np
a, b = test_input
[insert]
result = c
"""

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)


Prompt:
Problem:
Similar to this answer, I have a pair of 3D numpy arrays, a and b, and I want to sort the entries of b by the values of a. Unlike this answer, I want to sort only along one axis of the arrays, in decreasing order.
My naive reading of the numpy.argsort() documentation:
Returns
-------
index_array : ndarray, int
    Array of indices that sort `a` along the specified axis.
    In other words, ``a[index_array]`` yields a sorted `a`.
led me to believe that I could do my sort with the following code:
import numpy
print a
"""
[[[ 1.  1.  1.]
  [ 1.  1.  1.]
  [ 1.  1.  1.]]
 [[ 3.  3.  3.]
  [ 3.  2.  3.]
  [ 3.  3.  3.]]
 [[ 2.  2.  2.]
  [ 2.  3.  2.]
  [ 2.  2.  2.]]]
"""
b = numpy.arange(3*3*3).reshape((3, 3, 3))
print "b"
print b
"""
[[[ 0  1  2]
  [ 3  4  5]
  [ 6  7  8]]
 [[ 9 10 11]
  [12 13 14]
  [15 16 17]]
 [[18 19 20]
  [21 22 23]
  [24 25 26]]]
##This isnt' working how I'd like
sort_indices = numpy.argsort(a, axis=0)
c = b[sort_indices]
"""
Desired output:
[
 [[ 9 10 11]
  [12 22 14]
  [15 16 17]]
 [[18 19 20]
  [21 13 23]
  [24 25 26]]
 [[ 0  1  2]
  [ 3  4  5]
  [ 6  7  8]]]
"""
print "Desired shape of b[sort_indices]: (3, 3, 3)."
print "Actual shape of b[sort_indices]:"
print c.shape
"""
(3, 3, 3, 3, 3)
"""
What's the right way to do this?
A:
<code>
import numpy as np
a = np.random.rand(3, 3, 3)
b = np.arange(3*3*3).reshape((3, 3, 3))
</code>
c = ... # put solution in this variable
BEGIN SOLUTION
<code>

import numpy as np
import copy

def generate_test_case(test_case_id):
    def define_test_input(test_case_id):
        if test_case_id == 1:
            np.random.seed(42)
            a = np.random.rand(3, 3, 3)
            b = np.arange(3 * 3 * 3).reshape((3, 3, 3))
        return a, b

    def generate_ans(data):
        _a = data
        a, b = _a
        sort_indices = np.argsort(a, axis=0)[::-1, :, :]
        static_indices = np.indices(a.shape)
        c = b[sort_indices, static_indices[1], static_indices[2]]
        return c

    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):
    np.testing.assert_array_equal(result, ans)
    return 1

exec_context = r"""
import numpy as np
a, b = test_input
[insert]
result = c
"""

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)


Prompt:
Problem:
Similar to this answer, I have a pair of 3D numpy arrays, a and b, and I want to sort the matrices of b by the values of a. Unlike this answer, I want to sort the matrices according to their sum.
My naive reading of the numpy.argsort() documentation:
Returns
-------
index_array : ndarray, int
    Array of indices that sort `a` along the specified axis.
    In other words, ``a[index_array]`` yields a sorted `a`.
led me to believe that I could do my sort with the following code:
import numpy
print a
"""
[[[ 1.  1.  1.]
  [ 1.  1.  1.]
  [ 1.  1.  1.]]
 [[ 3.  3.  3.]
  [ 3.  2.  3.]
  [ 3.  3.  3.]]
 [[ 2.  2.  2.]
  [ 2.  3.  2.]
  [ 2.  2.  2.]]]
sum: 26 > 19 > 9
"""
b = numpy.arange(3*3*3).reshape((3, 3, 3))
print "b"
print b
"""
[[[ 0  1  2]
  [ 3  4  5]
  [ 6  7  8]]
 [[ 9 10 11]
  [12 13 14]
  [15 16 17]]
 [[18 19 20]
  [21 22 23]
  [24 25 26]]]

Desired output:
[[[ 0  1  2]
  [ 3  4  5]
  [ 6  7  8]]
 [[18 19 20]
  [21 22 23]
  [24 25 26]]
 [[ 9 10 11]
  [12 13 14]
  [15 16 17]]]

What's the right way to do this?
A:
<code>
import numpy as np
a = np.random.rand(3, 3, 3)
b = np.arange(3*3*3).reshape((3, 3, 3))
</code>
result = ... # put solution in this variable
BEGIN SOLUTION
<code>

index = np.argsort(a.sum(axis = (1, 2)))
result = b[index, :, :]

import numpy as np
import copy

def generate_test_case(test_case_id):
    def define_test_input(test_case_id):
        if test_case_id == 1:
            np.random.seed(42)
            a = np.random.rand(3, 3, 3)
            b = np.arange(3 * 3 * 3).reshape((3, 3, 3))
        return a, b

    def generate_ans(data):
        _a = data
        a, b = _a
        index = np.argsort(a.sum(axis=(1, 2)))
        result = b[index, :, :]
        return result

    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):
    np.testing.assert_array_equal(result, ans)
    return 1

exec_context = r"""
import numpy as np
a, b = 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)


Prompt:
Problem:

>>> arr = np.array([[1,2,3,4], [5,6,7,8], [9,10,11,12]])
>>> arr
array([[ 1,  2,  3,  4],
       [ 5,  6,  7,  8],
       [ 9, 10, 11, 12]])
I am deleting the 3rd column
array([[ 1,  2,  4],
       [ 5,  6,  8],
       [ 9, 10, 12]])
Are there any good way ?  Please consider this to be a novice question.
A:
<code>
import numpy as np
a = np.arange(12).reshape(3, 4)
</code>
a = ... # put solution in this variable
BEGIN SOLUTION
<code>

import numpy as np
import copy

def generate_test_case(test_case_id):
    def define_test_input(test_case_id):
        if test_case_id == 1:
            a = np.arange(12).reshape(3, 4)
        elif test_case_id == 2:
            a = np.ones((3, 3))
        return a

    def generate_ans(data):
        _a = data
        a = _a
        a = np.delete(a, 2, axis=1)
        return a

    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):
    np.testing.assert_array_equal(result, ans)
    return 1

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

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)


Prompt:
Problem:

>>> arr = np.array([[1,2,3,4], [5,6,7,8], [9,10,11,12]])
>>> arr
array([[ 1,  2,  3,  4],
       [ 5,  6,  7,  8],
       [ 9, 10, 11, 12]])
I am deleting the 3rd row
array([[ 1,  2,  3,  4],
       [ 5,  6,  7,  8]])
Are there any good way ?  Please consider this to be a novice question.

A:
<code>
import numpy as np
a = np.arange(12).reshape(3, 4)
</code>
a = ... # put solution in this variable
BEGIN SOLUTION
<code>

import numpy as np
import copy

def generate_test_case(test_case_id):
    def define_test_input(test_case_id):
        if test_case_id == 1:
            a = np.arange(12).reshape(3, 4)
        elif test_case_id == 2:
            a = np.ones((4, 4))
        return a

    def generate_ans(data):
        _a = data
        a = _a
        a = np.delete(a, 2, axis=0)
        return a

    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):
    np.testing.assert_array_equal(result, ans)
    return 1

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

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)


Prompt:
Problem:

>>> arr = np.array([[1,2,3,4], [5,6,7,8], [9,10,11,12]])
>>> arr
array([[ 1,  2,  3,  4],
       [ 5,  6,  7,  8],
       [ 9, 10, 11, 12]])
I am deleting the 1st and 3rd column
array([[ 2,  4],
       [ 6,  8],
       [ 10, 12]])
Are there any good way ? Please consider this to be a novice question.
A:
<code>
import numpy as np
a = np.arange(12).reshape(3, 4)
</code>
a = ... # put solution in this variable
BEGIN SOLUTION
<code>

temp = np.array([0, 2])
a = np.delete(a, temp, axis = 1)

import numpy as np
import copy

def generate_test_case(test_case_id):
    def define_test_input(test_case_id):
        if test_case_id == 1:
            a = np.arange(12).reshape(3, 4)
        elif test_case_id == 2:
            a = np.ones((6, 6))
        return a

    def generate_ans(data):
        _a = data
        a = _a
        temp = np.array([0, 2])
        a = np.delete(a, temp, axis=1)
        return a

    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):
    np.testing.assert_array_equal(result, ans)
    return 1

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

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)


Prompt:
Problem:

>>> arr = np.array([[1,2,3,4], [5,6,7,8], [9,10,11,12]])
>>> del_col = [1, 2, 4, 5]
>>> arr
array([[ 1,  2,  3,  4],
       [ 5,  6,  7,  8],
       [ 9, 10, 11, 12]])
I am deleting some columns(in this example, 1st, 2nd and 4th)
def_col = np.array([1, 2, 4, 5])
array([[ 3],
       [ 7],
       [ 11]])
Note that del_col might contain out-of-bound indices, so we should ignore them.
Are there any good way ? Please consider this to be a novice question.
A:
<code>
import numpy as np
a = np.arange(12).reshape(3, 4)
del_col = np.array([1, 2, 4, 5])
</code>
result = ... # put solution in this variable
BEGIN SOLUTION
<code>

mask = (del_col <= a.shape[1])
del_col = del_col[mask] - 1
result = np.delete(a, del_col, axis=1)

import numpy as np
import copy

def generate_test_case(test_case_id):
    def define_test_input(test_case_id):
        if test_case_id == 1:
            a = np.arange(12).reshape(3, 4)
            del_col = np.array([1, 2, 4, 5])
        elif test_case_id == 2:
            np.random.seed(42)
            a = np.random.rand(np.random.randint(5, 10), np.random.randint(6, 10))
            del_col = np.random.randint(0, 15, (4,))
        return a, del_col

    def generate_ans(data):
        _a = data
        a, del_col = _a
        mask = del_col <= a.shape[1]
        del_col = del_col[mask] - 1
        result = np.delete(a, del_col, axis=1)
        return result

    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):
    np.testing.assert_array_equal(result, ans)
    return 1

exec_context = r"""
import numpy as np
a, del_col = 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)


Prompt:
Problem:
Lists have a very simple method to insert elements:
a = [1,2,3,4]
a.insert(2,66)
print a
[1, 2, 66, 3, 4]
For a numpy array I could do:
a = np.asarray([1,2,3,4])
a_l = a.tolist()
a_l.insert(2,66)
a = np.asarray(a_l)
print a
[1 2 66 3 4]
but this is very convoluted.
Is there an insert equivalent for numpy arrays?
A:
<code>
import numpy as np
a = np.asarray([1,2,3,4])
pos = 2
element = 66
</code>
a = ... # put solution in this variable
BEGIN SOLUTION
<code>

import numpy as np
import pandas as pd
import copy
import tokenize, io

def generate_test_case(test_case_id):
    def define_test_input(test_case_id):
        if test_case_id == 1:
            a = np.asarray([1, 2, 3, 4])
            pos = 2
            element = 66
        elif test_case_id == 2:
            np.random.seed(42)
            a = np.random.rand(100)
            pos = np.random.randint(0, 99)
            element = 66
        return a, pos, element

    def generate_ans(data):
        _a = data
        a, pos, element = _a
        a = np.insert(a, pos, element)
        return a

    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):
    np.testing.assert_array_equal(result, ans)
    return 1

exec_context = r"""
import numpy as np
a, pos, element = test_input
[insert]
result = a
"""

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 test_string(solution: str):
    tokens = []
    for token in tokenize.tokenize(io.BytesIO(solution.encode("utf-8")).readline):
        tokens.append(token.string)
    assert "insert" in tokens


Prompt:
Problem:
Lists have a very simple method to insert elements:
a = [1,2,3,4]
a.insert(2,66)
print a
[1, 2, 66, 3, 4]
However, I’m confused about how to insert a row into an 2-dimensional array. e.g. changing
array([[1,2],[3,4]])
into
array([[1,2],[3,5],[3,4]])
A:
<code>
import numpy as np
a = np.array([[1,2],[3,4]])

pos = 1
element = [3,5]
</code>
a = ... # put solution in this variable
BEGIN SOLUTION
<code>

a = np.insert(a, pos, element, axis = 0)

import numpy as np
import pandas as pd
import copy
import tokenize, io

def generate_test_case(test_case_id):
    def define_test_input(test_case_id):
        if test_case_id == 1:
            a = np.array([[1, 2], [3, 4]])
            pos = 1
            element = [3, 5]
        elif test_case_id == 2:
            np.random.seed(42)
            a = np.random.rand(100, 10)
            pos = np.random.randint(0, 99)
            element = np.random.rand(10)
        return a, pos, element

    def generate_ans(data):
        _a = data
        a, pos, element = _a
        a = np.insert(a, pos, element, axis=0)
        return a

    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):
    np.testing.assert_array_equal(result, ans)
    return 1

exec_context = r"""
import numpy as np
a, pos, element = test_input
[insert]
result = a
"""

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 test_string(solution: str):
    tokens = []
    for token in tokenize.tokenize(io.BytesIO(solution.encode("utf-8")).readline):
        tokens.append(token.string)
    assert "insert" in tokens


Prompt:
Problem:
Lists have a very simple method to insert elements:
a = [1,2,3,4]
a.insert(2,66)
print a
[1, 2, 66, 3, 4]
For a numpy array I could do:
a = np.asarray([1,2,3,4])
a_l = a.tolist()
a_l.insert(2,66)
a = np.asarray(a_l)
print a
[1 2 66 3 4]
but this is very convoluted.
Is there an insert equivalent for numpy arrays?
A:
<code>
import numpy as np
example_a = np.asarray([1,2,3,4])
def f(a = example_a, pos=2, element = 66):
    # return the solution in this function
    # a = f(a, pos=2, element = 66)
    ### BEGIN SOLUTION

    a = np.insert(a, pos, element)

    return a

import numpy as np
import pandas as pd
import copy
import tokenize, io

def generate_test_case(test_case_id):
    def define_test_input(test_case_id):
        if test_case_id == 1:
            a = np.asarray([1, 2, 3, 4])
            pos = 2
            element = 66
        elif test_case_id == 2:
            np.random.seed(42)
            a = np.random.rand(100)
            pos = np.random.randint(0, 99)
            element = 66
        return a, pos, element

    def generate_ans(data):
        _a = data
        a, pos, element = _a
        a = np.insert(a, pos, element)
        return a

    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):
    np.testing.assert_array_equal(result, ans)
    return 1

exec_context = r"""
import numpy as np
a, pos, element = test_input
def f(a, pos=2, element = 66):
[insert]
result = f(a, pos, element)
"""

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 test_string(solution: str):
    tokens = []
    for token in tokenize.tokenize(io.BytesIO(solution.encode("utf-8")).readline):
        tokens.append(token.string)
    assert "insert" in tokens


Prompt:
Problem:
Lists have a very simple method to insert elements:
a = [1,2,3,4]
a.insert(2,66)
print a
[1, 2, 66, 3, 4]
However, I’m confused about how to insert multiple rows into an 2-dimensional array. Meanwhile, I want the inserted rows located in given indices in a. e.g.
a = array([[1,2],[3,4]])
element = array([[3, 5], [6, 6]])
pos = [1, 2]
array([[1,2],[3,5],[6,6], [3,4]])
Note that the given indices(pos) are monotonically increasing.
A:
<code>
import numpy as np
a = np.array([[1,2],[3,4]])
pos = [1, 2]
element = np.array([[3, 5], [6, 6]])
</code>
a = ... # put solution in this variable
BEGIN SOLUTION
<code>

import numpy as np
import pandas as pd
import copy
import tokenize, io

def generate_test_case(test_case_id):
    def define_test_input(test_case_id):
        if test_case_id == 1:
            a = np.array([[1, 2], [3, 4]])
            pos = [1, 2]
            element = np.array([[3, 5], [6, 6]])
        elif test_case_id == 2:
            np.random.seed(42)
            a = np.random.rand(100, 10)
            pos = sorted(np.random.randint(0, 99, (5,)))
            element = np.random.rand(5, 10)
        return a, pos, element

    def generate_ans(data):
        _a = data
        a, pos, element = _a
        pos = np.array(pos) - np.arange(len(element))
        a = np.insert(a, pos, element, axis=0)
        return a

    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):
    np.testing.assert_array_equal(result, ans)
    return 1

exec_context = r"""
import numpy as np
a, pos, element = test_input
[insert]
result = a
"""

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 test_string(solution: str):
    tokens = []
    for token in tokenize.tokenize(io.BytesIO(solution.encode("utf-8")).readline):
        tokens.append(token.string)
    assert "insert" in tokens


Prompt:
Problem:
I have a numpy array of different numpy arrays and I want to make a deep copy of the arrays. I found out the following:
import numpy as np
pairs = [(2, 3), (3, 4), (4, 5)]
array_of_arrays = np.array([np.arange(a*b).reshape(a,b) for (a, b) in pairs])
a = array_of_arrays[:] # Does not work
b = array_of_arrays[:][:] # Does not work
c = np.array(array_of_arrays, copy=True) # Does not work
Is for-loop the best way to do this? Is there a deep copy function I missed? And what is the best way to interact with each element in this array of different sized arrays?
A:
<code>
import numpy as np
pairs = [(2, 3), (3, 4), (4, 5)]
array_of_arrays = np.array([np.arange(a*b).reshape(a,b) for (a, b) in pairs])
</code>
result = ... # put solution in this variable
BEGIN SOLUTION
<code>

import copy
result = copy.deepcopy(array_of_arrays)

import numpy as np
import pandas as pd
import copy

def generate_test_case(test_case_id):
    def define_test_input(test_case_id):
        if test_case_id == 1:
            pairs = [(2, 3), (3, 4), (4, 5)]
            array_of_arrays = np.array(
                [np.arange(a * b).reshape(a, b) for (a, b) in pairs], dtype=object
            )
        return pairs, array_of_arrays

    def generate_ans(data):
        _a = data
        pairs, array_of_arrays = _a
        return array_of_arrays

    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):
    assert id(result) != id(ans)
    for arr1, arr2 in zip(result, ans):
        assert id(arr1) != id(arr2)
        np.testing.assert_array_equal(arr1, arr2)
    return 1

exec_context = r"""
import numpy as np
pairs, array_of_arrays = 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)


Prompt:
Problem:
In numpy, is there a nice idiomatic way of testing if all rows are equal in a 2d array?
I can do something like
np.all([np.array_equal(a[0], a[i]) for i in xrange(1,len(a))])
This seems to mix python lists with numpy arrays which is ugly and presumably also slow.
Is there a nicer/neater way?
A:
<code>
import numpy as np
a = np.repeat(np.arange(1, 6).reshape(1, -1), 3, axis = 0)
</code>
result = ... # put solution in this variable
BEGIN SOLUTION
<code>

import numpy as np
import pandas as pd
import copy
import tokenize, io

def generate_test_case(test_case_id):
    def define_test_input(test_case_id):
        if test_case_id == 1:
            a = np.repeat(np.arange(1, 6).reshape(1, -1), 3, axis=0)
        elif test_case_id == 2:
            np.random.seed(42)
            a = np.random.rand(3, 4)
        elif test_case_id == 3:
            a = np.array([[1, 1, 1], [2, 2, 2], [3, 3, 3]])
        elif test_case_id == 4:
            a = np.array([[1, 1, 1], [1, 1, 1], [1, 1, 3]])
        elif test_case_id == 5:
            a = np.array([[1, 1, 1], [2, 2, 1], [1, 1, 1]])
        return a

    def generate_ans(data):
        _a = data
        a = _a
        result = np.isclose(a, a[0], atol=0).all()
        return result

    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):
    assert result == ans
    return 1

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

def test_execution(solution: str):
    code = exec_context.replace("[insert]", solution)
    for i in range(5):
        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 test_string(solution: str):
    tokens = []
    for token in tokenize.tokenize(io.BytesIO(solution.encode("utf-8")).readline):
        tokens.append(token.string)
    assert "while" not in tokens and "for" not in tokens


Prompt:
Problem:
In numpy, is there a nice idiomatic way of testing if all columns are equal in a 2d array?
I can do something like
np.all([np.array_equal(a[0], a[i]) for i in xrange(1,len(a))])
This seems to mix python lists with numpy arrays which is ugly and presumably also slow.
Is there a nicer/neater way?
A:
<code>
import numpy as np
a = np.repeat(np.arange(1, 6).reshape(-1, 1), 3, axis = 1)
</code>
result = ... # put solution in this variable
BEGIN SOLUTION
<code>

result =np.isclose(a, a[:, 0].reshape(-1, 1), atol=0).all()

import numpy as np
import pandas as pd
import copy
import tokenize, io

def generate_test_case(test_case_id):
    def define_test_input(test_case_id):
        if test_case_id == 1:
            a = np.repeat(np.arange(1, 6).reshape(1, -1), 3, axis=0)
        elif test_case_id == 2:
            np.random.seed(42)
            a = np.random.rand(3, 4)
        elif test_case_id == 3:
            a = np.array([[1, 1, 1], [2, 2, 2], [3, 3, 3]])
        elif test_case_id == 4:
            a = np.array([[1, 1, 1], [1, 1, 2], [1, 1, 3]])
        elif test_case_id == 5:
            a = np.array([[1, 1, 1], [2, 2, 1], [3, 3, 1]])
        return a

    def generate_ans(data):
        _a = data
        a = _a
        result = np.isclose(a, a[:, 0].reshape(-1, 1), atol=0).all()
        return result

    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):
    assert result == ans
    return 1

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

def test_execution(solution: str):
    code = exec_context.replace("[insert]", solution)
    for i in range(5):
        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 test_string(solution: str):
    tokens = []
    for token in tokenize.tokenize(io.BytesIO(solution.encode("utf-8")).readline):
        tokens.append(token.string)
    assert "while" not in tokens and "for" not in tokens


Prompt:
Problem:
In numpy, is there a nice idiomatic way of testing if all rows are equal in a 2d array?
I can do something like
np.all([np.array_equal(a[0], a[i]) for i in xrange(1,len(a))])
This seems to mix python lists with numpy arrays which is ugly and presumably also slow.
Is there a nicer/neater way?
A:
<code>
import numpy as np
example_a = np.repeat(np.arange(1, 6).reshape(1, -1), 3, axis = 0)
def f(a = example_a):
    # return the solution in this function
    # result = f(a)
    ### BEGIN SOLUTION

import numpy as np
import pandas as pd
import copy
import tokenize, io

def generate_test_case(test_case_id):
    def define_test_input(test_case_id):
        if test_case_id == 1:
            a = np.repeat(np.arange(1, 6).reshape(1, -1), 3, axis=0)
        elif test_case_id == 2:
            np.random.seed(42)
            a = np.random.rand(3, 4)
        elif test_case_id == 3:
            a = np.array([[1, 1, 1], [2, 2, 2], [3, 3, 3]])
        elif test_case_id == 4:
            a = np.array([[1, 1, 1], [1, 1, 1], [1, 1, 3]])
        elif test_case_id == 5:
            a = np.array([[1, 1, 1], [2, 2, 1], [1, 1, 1]])
        return a

    def generate_ans(data):
        _a = data
        a = _a
        result = np.isclose(a, a[0], atol=0).all()
        return result

    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):
    assert result == ans
    return 1

exec_context = r"""
import numpy as np
a = test_input
def f(a):
[insert]
result = f(a)
"""

def test_execution(solution: str):
    code = exec_context.replace("[insert]", solution)
    for i in range(5):
        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 test_string(solution: str):
    tokens = []
    for token in tokenize.tokenize(io.BytesIO(solution.encode("utf-8")).readline):
        tokens.append(token.string)
    assert "while" not in tokens and "for" not in tokens


Prompt:
Problem:
SciPy has three methods for doing 1D integrals over samples (trapz, simps, and romb) and one way to do a 2D integral over a function (dblquad), but it doesn't seem to have methods for doing a 2D integral over samples -- even ones on a rectangular grid.
The closest thing I see is scipy.interpolate.RectBivariateSpline.integral -- you can create a RectBivariateSpline from data on a rectangular grid and then integrate it. However, that isn't terribly fast.
I want something more accurate than the rectangle method (i.e. just summing everything up). I could, say, use a 2D Simpson's rule by making an array with the correct weights, multiplying that by the array I want to integrate, and then summing up the result.
However, I don't want to reinvent the wheel if there's already something better out there. Is there?
For instance, I want to do 2D integral over (cosx)^4 + (siny)^2, how can I do it? Perhaps using Simpson rule?
A:
<code>
import numpy as np
x = np.linspace(0, 1, 20)
y = np.linspace(0, 1, 30)
</code>
result = ... # put solution in this variable
BEGIN SOLUTION
<code>

from scipy.integrate import simpson
z = np.cos(x[:,None])**4 + np.sin(y)**2
result = simpson(simpson(z, y), x)

import numpy as np
import copy
import scipy
from scipy.integrate import simpson

def generate_test_case(test_case_id):

    def define_test_input(test_case_id):
        if test_case_id == 1:
            x = np.linspace(0, 1, 20)
            y = np.linspace(0, 1, 30)
        elif test_case_id == 2:
            x = np.linspace(3, 5, 30)
            y = np.linspace(0, 1, 20)
        return x, y

    def generate_ans(data):
        _a = data
        x, y = _a
        z = np.cos(x[:, None]) ** 4 + np.sin(y) ** 2
        result = simpson(simpson(z, y), x)
        return result

    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):
    np.testing.assert_allclose(result, ans)
    return 1

exec_context = r"""
import numpy as np
x, y = 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)


Prompt:
Problem:
SciPy has three methods for doing 1D integrals over samples (trapz, simps, and romb) and one way to do a 2D integral over a function (dblquad), but it doesn't seem to have methods for doing a 2D integral over samples -- even ones on a rectangular grid.
The closest thing I see is scipy.interpolate.RectBivariateSpline.integral -- you can create a RectBivariateSpline from data on a rectangular grid and then integrate it. However, that isn't terribly fast.
I want something more accurate than the rectangle method (i.e. just summing everything up). I could, say, use a 2D Simpson's rule by making an array with the correct weights, multiplying that by the array I want to integrate, and then summing up the result.
However, I don't want to reinvent the wheel if there's already something better out there. Is there?
For instance, I want to do 2D integral over (cosx)^4 + (siny)^2, how can I do it? Perhaps using Simpson rule?
A:
<code>
import numpy as np
example_x = np.linspace(0, 1, 20)
example_y = np.linspace(0, 1, 30)
def f(x = example_x, y = example_y):
    # return the solution in this function
    # result = f(x, y)
    ### BEGIN SOLUTION

    from scipy.integrate import simpson
    z = np.cos(x[:,None])**4 + np.sin(y)**2
    result = simpson(simpson(z, y), x)

    return result

import numpy as np
import copy
import scipy
from scipy.integrate import simpson

def generate_test_case(test_case_id):

    def define_test_input(test_case_id):
        if test_case_id == 1:
            x = np.linspace(0, 1, 20)
            y = np.linspace(0, 1, 30)
        elif test_case_id == 2:
            x = np.linspace(3, 5, 30)
            y = np.linspace(0, 1, 20)
        return x, y

    def generate_ans(data):
        _a = data
        x, y = _a
        z = np.cos(x[:, None]) ** 4 + np.sin(y) ** 2
        result = simpson(simpson(z, y), x)
        return result

    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):
    np.testing.assert_allclose(result, ans)
    return 1

exec_context = r"""
import numpy as np
x, y = test_input
def f(x, y):
[insert]
result = f(x, y)
"""

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)


Prompt:
Problem:
What is the equivalent of R's ecdf(x)(x) function in Python, in either numpy or scipy? Is ecdf(x)(x) basically the same as:
import numpy as np
def ecdf(x):
  # normalize X to sum to 1
  x = x / np.sum(x)
  return np.cumsum(x)
or is something else required?
By default R's ecdf will return function values of elements in x in increasing order, and I want to get that in Python.
A:
<code>
import numpy as np
grades = np.array((93.5,93,60.8,94.5,82,87.5,91.5,99.5,86,93.5,92.5,78,76,69,94.5,
          89.5,92.8,78,65.5,98,98.5,92.3,95.5,76,91,95,61))
</code>
result = ... # put solution in this variable
BEGIN SOLUTION
<code>

def ecdf_result(x):
    xs = np.sort(x)
    ys = np.arange(1, len(xs)+1)/float(len(xs))
    return ys
result = ecdf_result(grades)

import numpy as np
import pandas as pd
import copy

def generate_test_case(test_case_id):
    def define_test_input(test_case_id):
        if test_case_id == 1:
            grades = np.array(
                (
                    93.5,
                    93,
                    60.8,
                    94.5,
                    82,
                    87.5,
                    91.5,
                    99.5,
                    86,
                    93.5,
                    92.5,
                    78,
                    76,
                    69,
                    94.5,
                    89.5,
                    92.8,
                    78,
                    65.5,
                    98,
                    98.5,
                    92.3,
                    95.5,
                    76,
                    91,
                    95,
                    61,
                )
            )
        elif test_case_id == 2:
            np.random.seed(42)
            grades = (np.random.rand(50) - 0.5) * 100
        return grades

    def generate_ans(data):
        _a = data
        grades = _a

        def ecdf_result(x):
            xs = np.sort(x)
            ys = np.arange(1, len(xs) + 1) / float(len(xs))
            return ys

        result = ecdf_result(grades)
        return result

    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):
    assert np.allclose(result, ans)
    return 1

exec_context = r"""
import numpy as np
grades = 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)


Prompt:
Problem:
What is the equivalent of R's ecdf(x)(x) function in Python, in either numpy or scipy? Is ecdf(x)(x) basically the same as:
import numpy as np
def ecdf(x):
  # normalize X to sum to 1
  x = x / np.sum(x)
  return np.cumsum(x)
or is something else required?
What I want to do is to apply the generated ECDF function to an eval array to gets corresponding values for elements in it.
A:
<code>
import numpy as np
grades = np.array((93.5,93,60.8,94.5,82,87.5,91.5,99.5,86,93.5,92.5,78,76,69,94.5,
          89.5,92.8,78,65.5,98,98.5,92.3,95.5,76,91,95,61))
eval = np.array([88, 87, 62])
</code>
result = ... # put solution in this variable
BEGIN SOLUTION
<code>

def ecdf_result(x):
    xs = np.sort(x)
    ys = np.arange(1, len(xs)+1)/float(len(xs))
    return xs, ys
resultx, resulty = ecdf_result(grades)
result = np.zeros_like(eval, dtype=float)
for i, element in enumerate(eval):
    if element < resultx[0]:
        result[i] = 0
    elif element >= resultx[-1]:
        result[i] = 1
    else:
        result[i] = resulty[(resultx > element).argmax()-1]

import numpy as np
import pandas as pd
import copy

def generate_test_case(test_case_id):
    def define_test_input(test_case_id):
        if test_case_id == 1:
            grades = np.array(
                (
                    93.5,
                    93,
                    60.8,
                    94.5,
                    82,
                    87.5,
                    91.5,
                    99.5,
                    86,
                    93.5,
                    92.5,
                    78,
                    76,
                    69,
                    94.5,
                    89.5,
                    92.8,
                    78,
                    65.5,
                    98,
                    98.5,
                    92.3,
                    95.5,
                    76,
                    91,
                    95,
                    61,
                )
            )
            eval = np.array([88, 87, 62])
        elif test_case_id == 2:
            np.random.seed(42)
            grades = (np.random.rand(50) - 0.5) * 100
            eval = np.random.randint(10, 90, (5,))
        return grades, eval

    def generate_ans(data):
        _a = data
        grades, eval = _a

        def ecdf_result(x):
            xs = np.sort(x)
            ys = np.arange(1, len(xs) + 1) / float(len(xs))
            return xs, ys

        resultx, resulty = ecdf_result(grades)
        result = np.zeros_like(eval, dtype=float)
        for i, element in enumerate(eval):
            if element < resultx[0]:
                result[i] = 0
            elif element >= resultx[-1]:
                result[i] = 1
            else:
                result[i] = resulty[(resultx > element).argmax() - 1]
        return result

    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):
    assert np.allclose(result, ans)
    return 1

exec_context = r"""
import numpy as np
grades, eval = 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)


Prompt:
Problem:
What is the equivalent of R's ecdf(x)(x) function in Python, in either numpy or scipy? Is ecdf(x)(x) basically the same as:
import numpy as np
def ecdf(x):
  # normalize X to sum to 1
  x = x / np.sum(x)
  return np.cumsum(x)
or is something else required?
Further, I want to compute the longest interval [low, high) that satisfies ECDF(x) < threshold for any x in [low, high). Note that low, high are elements of original array.
A:
<code>
import numpy as np
grades = np.array((93.5,93,60.8,94.5,82,87.5,91.5,99.5,86,93.5,92.5,78,76,69,94.5,
          89.5,92.8,78,65.5,98,98.5,92.3,95.5,76,91,95,61))
threshold = 0.5
</code>
low, high = ... # put solution in these variables
BEGIN SOLUTION
<code>

def ecdf_result(x):
    xs = np.sort(x)
    ys = np.arange(1, len(xs)+1)/float(len(xs))
    return xs, ys
resultx, resulty = ecdf_result(grades)
t = (resulty > threshold).argmax()
low = resultx[0]
high = resultx[t]

import numpy as np
import pandas as pd
import copy

def generate_test_case(test_case_id):
    def define_test_input(test_case_id):
        if test_case_id == 1:
            grades = np.array(
                (
                    93.5,
                    93,
                    60.8,
                    94.5,
                    82,
                    87.5,
                    91.5,
                    99.5,
                    86,
                    93.5,
                    92.5,
                    78,
                    76,
                    69,
                    94.5,
                    89.5,
                    92.8,
                    78,
                    65.5,
                    98,
                    98.5,
                    92.3,
                    95.5,
                    76,
                    91,
                    95,
                    61,
                )
            )
            threshold = 0.5
        elif test_case_id == 2:
            np.random.seed(42)
            grades = (np.random.rand(50) - 0.5) * 100
            threshold = 0.6
        return grades, threshold

    def generate_ans(data):
        _a = data
        grades, threshold = _a

        def ecdf_result(x):
            xs = np.sort(x)
            ys = np.arange(1, len(xs) + 1) / float(len(xs))
            return xs, ys

        resultx, resulty = ecdf_result(grades)
        t = (resulty > threshold).argmax()
        low = resultx[0]
        high = resultx[t]
        return [low, high]

    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):
    np.testing.assert_allclose(result, ans)
    return 1

exec_context = r"""
import numpy as np
grades, threshold = test_input
[insert]
result = [low, high]
"""

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)


Prompt:
Problem:
I want to generate a random array of size N which only contains 0 and 1, I want my array to have some ratio between 0 and 1. For example, 90% of the array be 1 and the remaining 10% be 0 (I want this 90% to be random along with the whole array).
right now I have:
randomLabel = np.random.randint(2, size=numbers)
But I can't control the ratio between 0 and 1.
A:
<code>
import numpy as np
one_ratio = 0.9
size = 1000
</code>
nums = ... # put solution in this variable
BEGIN SOLUTION
<code>

nums = np.ones(size)
nums[:int(size*(1-one_ratio))] = 0
np.random.shuffle(nums)

import numpy as np
import pandas as pd
import copy
import tokenize, io

def generate_test_case(test_case_id):
    def define_test_input(test_case_id):
        if test_case_id == 1:
            one_ratio = 0.9
            size = 1000
        elif test_case_id == 2:
            size = 100
            one_ratio = 0.8
        return size, one_ratio

    def generate_ans(data):
        _a = data
        return _a

    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):
    size, one_ratio = ans
    assert result.shape == (size,)
    assert abs(len(np.where(result == 1)[0]) - size * one_ratio) / size <= 0.05
    assert abs(len(np.where(result == 0)[0]) - size * (1 - one_ratio)) / size <= 0.05
    return 1

exec_context = r"""
import numpy as np
size, one_ratio = test_input
[insert]
result = nums
"""

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 test_string(solution: str):
    tokens = []
    for token in tokenize.tokenize(io.BytesIO(solution.encode("utf-8")).readline):
        tokens.append(token.string)
    assert "random" in tokens


Prompt:
Problem:
How do I convert a torch tensor to numpy?
A:
<code>
import torch
import numpy as np
a = torch.ones(5)
</code>
a_np = ... # put solution in this variable
BEGIN SOLUTION
<code>

import numpy as np
import pandas as pd
import torch
import copy

def generate_test_case(test_case_id):
    def define_test_input(test_case_id):
        if test_case_id == 1:
            a = torch.ones(5)
        elif test_case_id == 2:
            a = torch.tensor([1, 1, 4, 5, 1, 4])
        return a

    def generate_ans(data):
        _a = data
        a = _a
        a_np = a.numpy()
        return a_np

    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):
    assert type(result) == np.ndarray
    np.testing.assert_array_equal(result, ans)
    return 1

exec_context = r"""
import torch
import numpy as np
a = test_input
[insert]
result = a_np
"""

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)


Prompt:
Problem:
How do I convert a numpy array to pytorch tensor?
A:
<code>
import torch
import numpy as np
a = np.ones(5)
</code>
a_pt = ... # put solution in this variable
BEGIN SOLUTION
<code>

import numpy as np
import pandas as pd
import torch
import copy

def generate_test_case(test_case_id):
    def define_test_input(test_case_id):
        if test_case_id == 1:
            a = np.ones(5)
        elif test_case_id == 2:
            a = np.array([1, 1, 4, 5, 1, 4])
        return a

    def generate_ans(data):
        _a = data
        a = _a
        a_pt = torch.Tensor(a)
        return a_pt

    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):
    assert type(result) == torch.Tensor
    torch.testing.assert_close(result, ans, check_dtype=False)
    return 1

exec_context = r"""
import torch
import numpy as np
a = test_input
[insert]
result = a_pt
"""

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)


Prompt:
Problem:
How do I convert a tensorflow tensor to numpy?
A:
<code>
import tensorflow as tf
import numpy as np
a = tf.ones([2,3,4])
</code>
a_np = ... # put solution in this variable
BEGIN SOLUTION
<code>

import numpy as np
import pandas as pd
import tensorflow as tf
import copy

def generate_test_case(test_case_id):
    def define_test_input(test_case_id):
        if test_case_id == 1:
            a = tf.ones([2, 3, 4])
        elif test_case_id == 2:
            a = tf.zeros([3, 4])
        return a

    def generate_ans(data):
        _a = data
        a = _a
        a_np = a.numpy()
        return a_np

    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):
    assert type(result) == np.ndarray
    np.testing.assert_array_equal(result, ans)
    return 1

exec_context = r"""
import tensorflow as tf
import numpy as np
a = test_input
[insert]
result = a_np
"""

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)


Prompt:
Problem:
How do I convert a numpy array to tensorflow tensor?
A:
<code>
import tensorflow as tf
import numpy as np
a = np.ones([2,3,4])
</code>
a_tf = ... # put solution in this variable
BEGIN SOLUTION
<code>

import numpy as np
import pandas as pd
import tensorflow as tf
import copy

def generate_test_case(test_case_id):
    def define_test_input(test_case_id):
        if test_case_id == 1:
            a = np.ones([2, 3, 4])
        elif test_case_id == 2:
            a = np.array([1, 1, 4, 5, 1, 4])
        return a

    def generate_ans(data):
        _a = data
        a = _a
        a_tf = tf.convert_to_tensor(a)
        return a_tf

    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):
    def tensor_equal(a, b):
        if type(a) != type(b):
            return False
        if isinstance(a, type(tf.constant([]))) is not True:
            if isinstance(a, type(tf.Variable([]))) is not True:
                return False
        if a.shape != b.shape:
            return False
        if a.dtype != tf.float32:
            a = tf.cast(a, tf.float32)
        if b.dtype != tf.float32:
            b = tf.cast(b, tf.float32)
        if not tf.reduce_min(tf.cast(a == b, dtype=tf.int32)):
            return False
        return True

    assert tensor_equal(result, ans)
    return 1

exec_context = r"""
import tensorflow as tf
import numpy as np
a = test_input
[insert]
result = a_tf
"""

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)


Prompt:
Problem:
I'm sorry in advance if this is a duplicated question, I looked for this information but still couldn't find it.
Is it possible to get a numpy array (or python list) filled with the indexes of the elements in decreasing order?
For instance, the array:
a = array([4, 1, 0, 8, 5, 2])
The indexes of the elements in decreasing order would give :
8 --> 3
5 --> 4
4 --> 0
2 --> 5
1 --> 1
0 --> 2
result = [3, 4, 0, 5, 1, 2]
Thanks in advance!
A:
<code>
import numpy as np
a = np.array([4, 1, 0, 8, 5, 2])
</code>
result = ... # put solution in this variable
BEGIN SOLUTION
<code>

import numpy as np
import copy

def generate_test_case(test_case_id):
    def define_test_input(test_case_id):
        if test_case_id == 1:
            a = np.array([4, 1, 0, 8, 5, 2])
        elif test_case_id == 2:
            np.random.seed(42)
            a = (np.random.rand(100) - 0.5) * 100
        return a

    def generate_ans(data):
        _a = data
        a = _a
        result = np.argsort(a)[::-1][: len(a)]
        return result

    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):
    np.testing.assert_array_equal(result, ans)
    return 1

exec_context = r"""
import numpy as np
a = 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)


Prompt:
Problem:
I'm sorry in advance if this is a duplicated question, I looked for this information but still couldn't find it.
Is it possible to get a numpy array (or python list) filled with the indexes of the elements in increasing order?
For instance, the array:
a = array([4, 1, 0, 8, 5, 2])
The indexes of the elements in increasing order would give :
0 --> 2
1 --> 1
2 --> 5
4 --> 0
5 --> 4
8 --> 3
result = [2,1,5,0,4,3]
Thanks in advance!
A:
<code>
import numpy as np
a = np.array([4, 1, 0, 8, 5, 2])
</code>
result = ... # put solution in this variable
BEGIN SOLUTION
<code>

import numpy as np
import copy

def generate_test_case(test_case_id):
    def define_test_input(test_case_id):
        if test_case_id == 1:
            a = np.array([4, 1, 0, 8, 5, 2])
        elif test_case_id == 2:
            np.random.seed(42)
            a = (np.random.rand(100) - 0.5) * 100
        return a

    def generate_ans(data):
        _a = data
        a = _a
        result = np.argsort(a)
        return result

    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):
    np.testing.assert_array_equal(result, ans)
    return 1

exec_context = r"""
import numpy as np
a = 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)


Prompt:
Problem:
I'm sorry in advance if this is a duplicated question, I looked for this information but still couldn't find it.
Is it possible to get a numpy array (or python list) filled with the indexes of the N biggest elements in decreasing order?
For instance, the array:
a = array([4, 1, 0, 8, 5, 2])
The indexes of the biggest elements in decreasing order would give (considering N = 3):
8 --> 3
5 --> 4
4 --> 0
result = [3, 4, 0]
Thanks in advance!
A:
<code>
import numpy as np
a = np.array([4, 1, 0, 8, 5, 2])
N = 3
</code>
result = ... # put solution in this variable
BEGIN SOLUTION
<code>

import numpy as np
import copy

def generate_test_case(test_case_id):
    def define_test_input(test_case_id):
        if test_case_id == 1:
            a = np.array([4, 1, 0, 8, 5, 2])
            N = 3
        elif test_case_id == 2:
            np.random.seed(42)
            a = (np.random.rand(100) - 0.5) * 100
            N = np.random.randint(1, 25)
        return a, N

    def generate_ans(data):
        _a = data
        a, N = _a
        result = np.argsort(a)[::-1][:N]
        return result

    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):
    np.testing.assert_array_equal(result, ans)
    return 1

exec_context = r"""
import numpy as np
a, N = 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)


Prompt:
Problem:

I want to raise a 2-dimensional numpy array, let's call it A, to the power of some number n, but I have thus far failed to find the function or operator to do that.
I'm aware that I could cast it to the matrix type and use the fact that then (similar to what would be the behaviour in Matlab), A**n does just what I want, (for array the same expression means elementwise exponentiation). Casting to matrix and back seems like a rather ugly workaround though.
Surely there must be a good way to perform that calculation while keeping the format to array?
A:
<code>
import numpy as np
A = np.arange(16).reshape(4, 4)
n = 5
</code>
result = ... # put solution in this variable
BEGIN SOLUTION
<code>

import numpy as np
import pandas as pd
import copy
import tokenize, io

def generate_test_case(test_case_id):
    def define_test_input(test_case_id):
        if test_case_id == 1:
            A = np.arange(16).reshape(4, 4)
            n = 5
        elif test_case_id == 2:
            np.random.seed(42)
            dim = np.random.randint(10, 15)
            A = np.random.rand(dim, dim)
            n = np.random.randint(3, 8)
        return A, n

    def generate_ans(data):
        _a = data
        A, n = _a
        result = np.linalg.matrix_power(A, n)
        return result

    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):
    assert type(result) == np.ndarray
    np.testing.assert_allclose(result, ans)
    return 1

exec_context = r"""
import numpy as np
A, n = 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 test_string(solution: str):
    tokens = []
    for token in tokenize.tokenize(io.BytesIO(solution.encode("utf-8")).readline):
        tokens.append(token.string)
    assert "matrix" not in tokens


Prompt:
Problem:
I have a 2-d numpy array as follows:
a = np.array([[1,5,9,13],
              [2,6,10,14],
              [3,7,11,15],
              [4,8,12,16]]
I want to extract it into patches of 2 by 2 sizes with out repeating the elements.
The answer should exactly be the same. This can be 3-d array or list with the same order of elements as below:
[[[1,5],
 [2,6]],
 [[3,7],
 [4,8]],
 [[9,13],
 [10,14]],
 [[11,15],
 [12,16]]]
How can do it easily?
In my real problem the size of a is (36, 72). I can not do it one by one. I want programmatic way of doing it.
A:
<code>
import numpy as np
a = np.array([[1,5,9,13],
              [2,6,10,14],
              [3,7,11,15],
              [4,8,12,16]])
</code>
result = ... # put solution in this variable
BEGIN SOLUTION
<code>

result = a.reshape(a.shape[0]//2, 2, a.shape[1]//2, 2).swapaxes(1, 2).transpose(1, 0, 2, 3).reshape(-1, 2, 2)

import numpy as np
import copy
import tokenize, io

def generate_test_case(test_case_id):
    def define_test_input(test_case_id):
        if test_case_id == 1:
            a = np.array(
                [[1, 5, 9, 13], [2, 6, 10, 14], [3, 7, 11, 15], [4, 8, 12, 16]]
            )
        elif test_case_id == 2:
            np.random.seed(42)
            a = np.random.rand(100, 200)
        return a

    def generate_ans(data):
        _a = data
        a = _a
        result = (
            a.reshape(a.shape[0] // 2, 2, a.shape[1] // 2, 2)
            .swapaxes(1, 2)
            .transpose(1, 0, 2, 3)
            .reshape(-1, 2, 2)
        )
        return result

    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):
    np.testing.assert_array_equal(result, ans)
    return 1

exec_context = r"""
import numpy as np
a = 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 test_string(solution: str):
    tokens = []
    for token in tokenize.tokenize(io.BytesIO(solution.encode("utf-8")).readline):
        tokens.append(token.string)
    assert "while" not in tokens and "for" not in tokens


Prompt:
Problem:
I have a 2-d numpy array as follows:
a = np.array([[1,5,9,13],
              [2,6,10,14],
              [3,7,11,15],
              [4,8,12,16]]
I want to extract it into patches of 2 by 2 sizes like sliding window.
The answer should exactly be the same. This can be 3-d array or list with the same order of elements as below:
[[[1,5],
 [2,6]],
 [[5,9],
 [6,10]],
 [[9,13],
 [10,14]],
 [[2,6],
 [3,7]],
 [[6,10],
 [7,11]],
 [[10,14],
 [11,15]],
 [[3,7],
 [4,8]],
 [[7,11],
 [8,12]],
 [[11,15],
 [12,16]]]
How can do it easily?
In my real problem the size of a is (36, 72). I can not do it one by one. I want programmatic way of doing it.
A:
<code>
import numpy as np
a = np.array([[1,5,9,13],
              [2,6,10,14],
              [3,7,11,15],
              [4,8,12,16]])
</code>
result = ... # put solution in this variable
BEGIN SOLUTION
<code>

result = np.lib.stride_tricks.sliding_window_view(a, window_shape=(2,2)).reshape(-1, 2, 2)

import numpy as np
import copy
import tokenize, io

def generate_test_case(test_case_id):
    def define_test_input(test_case_id):
        if test_case_id == 1:
            a = np.array(
                [[1, 5, 9, 13], [2, 6, 10, 14], [3, 7, 11, 15], [4, 8, 12, 16]]
            )
        elif test_case_id == 2:
            np.random.seed(42)
            a = np.random.rand(100, 200)
        return a

    def generate_ans(data):
        _a = data
        a = _a
        result = np.lib.stride_tricks.sliding_window_view(
            a, window_shape=(2, 2)
        ).reshape(-1, 2, 2)
        return result

    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):
    np.testing.assert_array_equal(result, ans)
    return 1

exec_context = r"""
import numpy as np
a = 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 test_string(solution: str):
    tokens = []
    for token in tokenize.tokenize(io.BytesIO(solution.encode("utf-8")).readline):
        tokens.append(token.string)
    assert "while" not in tokens and "for" not in tokens


Prompt:
Problem:
I have a 2-d numpy array as follows:
a = np.array([[1,5,9,13],
              [2,6,10,14],
              [3,7,11,15],
              [4,8,12,16]]
I want to extract it into patches of 2 by 2 sizes with out repeating the elements.
The answer should exactly be the same. This can be 3-d array or list with the same order of elements as below:
[[[1,5],
 [2,6]],
 [[9,13],
 [10,14]],
 [[3,7],
 [4,8]],
 [[11,15],
 [12,16]]]
How can do it easily?
In my real problem the size of a is (36, 72). I can not do it one by one. I want programmatic way of doing it.
A:
<code>
import numpy as np
a = np.array([[1,5,9,13],
              [2,6,10,14],
              [3,7,11,15],
              [4,8,12,16]])
</code>
result = ... # put solution in this variable
BEGIN SOLUTION
<code>

result = a.reshape(a.shape[0]//2, 2, a.shape[1]//2, 2).swapaxes(1, 2).reshape(-1, 2, 2)

import numpy as np
import copy
import tokenize, io

def generate_test_case(test_case_id):
    def define_test_input(test_case_id):
        if test_case_id == 1:
            a = np.array(
                [[1, 5, 9, 13], [2, 6, 10, 14], [3, 7, 11, 15], [4, 8, 12, 16]]
            )
        elif test_case_id == 2:
            np.random.seed(42)
            a = np.random.rand(100, 200)
        return a

    def generate_ans(data):
        _a = data
        a = _a
        result = (
            a.reshape(a.shape[0] // 2, 2, a.shape[1] // 2, 2)
            .swapaxes(1, 2)
            .reshape(-1, 2, 2)
        )
        return result

    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):
    np.testing.assert_array_equal(result, ans)
    return 1

exec_context = r"""
import numpy as np
a = 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 test_string(solution: str):
    tokens = []
    for token in tokenize.tokenize(io.BytesIO(solution.encode("utf-8")).readline):
        tokens.append(token.string)
    assert "while" not in tokens and "for" not in tokens


Prompt:
Problem:
I have a 2-d numpy array as follows:
a = np.array([[1,5,9,13,17],
              [2,6,10,14,18],
              [3,7,11,15,19],
              [4,8,12,16,20]]
I want to extract it into patches of 2 by 2 sizes with out repeating the elements. Pay attention that if the shape is indivisible by patch size, we would just ignore the rest row/column.
The answer should exactly be the same. This can be 3-d array or list with the same order of elements as below:
[[[1,5],
 [2,6]],
 [[9,13],
 [10,14]],
 [[3,7],
 [4,8]],
 [[11,15],
 [12,16]]]
How can do it easily?
In my real problem the size of a is (36, 73). I can not do it one by one. I want programmatic way of doing it.
A:
<code>
import numpy as np
a = np.array([[1,5,9,13,17],
              [2,6,10,14,18],
              [3,7,11,15,19],
              [4,8,12,16,20]])
patch_size = 2
</code>
result = ... # put solution in this variable
BEGIN SOLUTION
<code>

x = a[:a.shape[0] // patch_size * patch_size, :a.shape[1] // patch_size * patch_size]
result = x.reshape(x.shape[0]//patch_size, patch_size, x.shape[1]// patch_size, patch_size).swapaxes(1, 2). reshape(-1, patch_size, patch_size)

import numpy as np
import copy
import tokenize, io

def generate_test_case(test_case_id):
    def define_test_input(test_case_id):
        if test_case_id == 1:
            a = np.array(
                [
                    [1, 5, 9, 13, 17],
                    [2, 6, 10, 14, 18],
                    [3, 7, 11, 15, 19],
                    [4, 8, 12, 16, 20],
                ]
            )
            patch_size = 2
        elif test_case_id == 2:
            np.random.seed(42)
            a = np.random.rand(100, 200)
            patch_size = np.random.randint(4, 8)
        return a, patch_size

    def generate_ans(data):
        _a = data
        a, patch_size = _a
        x = a[
            : a.shape[0] // patch_size * patch_size,
            : a.shape[1] // patch_size * patch_size,
        ]
        result = (
            x.reshape(
                x.shape[0] // patch_size,
                patch_size,
                x.shape[1] // patch_size,
                patch_size,
            )
            .swapaxes(1, 2)
            .reshape(-1, patch_size, patch_size)
        )
        return result

    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):
    np.testing.assert_array_equal(result, ans)
    return 1

exec_context = r"""
import numpy as np
a, patch_size = 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 test_string(solution: str):
    tokens = []
    for token in tokenize.tokenize(io.BytesIO(solution.encode("utf-8")).readline):
        tokens.append(token.string)
    assert "while" not in tokens and "for" not in tokens


Prompt:
Problem:
I'm looking for a generic method to from the original big array from small arrays:
array([[[ 0,  1,  2],
        [ 6,  7,  8]],
       [[ 3,  4,  5],
        [ 9, 10, 11]],
       [[12, 13, 14],
        [18, 19, 20]],
       [[15, 16, 17],
        [21, 22, 23]]])
->
# result array's shape: (h = 4, w = 6)
array([[ 0,  1,  2,  3,  4,  5],
       [ 6,  7,  8,  9, 10, 11],
       [12, 13, 14, 15, 16, 17],
       [18, 19, 20, 21, 22, 23]])
I am currently developing a solution, will post it when it's done, would however like to see other (better) ways.
A:
<code>
import numpy as np
a = np.array([[[ 0,  1,  2],
        [ 6,  7,  8]],
       [[ 3,  4,  5],
        [ 9, 10, 11]],
       [[12, 13, 14],
        [18, 19, 20]],
       [[15, 16, 17],
        [21, 22, 23]]])
h = 4
w = 6
</code>
result = ... # put solution in this variable
BEGIN SOLUTION
<code>

n, nrows, ncols = a.shape
result = a.reshape(h//nrows, -1, nrows, ncols).swapaxes(1,2).reshape(h, w)

import numpy as np
import copy

def generate_test_case(test_case_id):
    def define_test_input(test_case_id):
        if test_case_id == 1:
            a = np.array(
                [
                    [[0, 1, 2], [6, 7, 8]],
                    [[3, 4, 5], [9, 10, 11]],
                    [[12, 13, 14], [18, 19, 20]],
                    [[15, 16, 17], [21, 22, 23]],
                ]
            )
            h = 4
            w = 6
        elif test_case_id == 2:
            np.random.seed(42)
            a = np.random.rand(6, 3, 4)
            h = 6
            w = 12
        return a, h, w

    def generate_ans(data):
        _a = data
        a, h, w = _a
        n, nrows, ncols = a.shape
        result = a.reshape(h // nrows, -1, nrows, ncols).swapaxes(1, 2).reshape(h, w)
        return result

    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):
    np.testing.assert_array_equal(result, ans)
    return 1

exec_context = r"""
import numpy as np
a, h, w = 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)


Prompt:
Problem:
I have a 2-d numpy array as follows:
a = np.array([[1,5,9,13,17],
              [2,6,10,14,18],
              [3,7,11,15,19],
              [4,8,12,16,20]]
I want to extract it into patches of 2 by 2 sizes with out repeating the elements. Pay attention that if the shape is indivisible by patch size, we would just ignore the rest row/column.
The answer should exactly be the same. This can be 3-d array or list with the same order of elements as below:
[[[1,5],
 [2,6]],
 [[3,7],
 [4,8]],
 [[9,13],
 [10,14]],
 [[11,15],
 [12,16]]]
How can do it easily?
In my real problem the size of a is (36, 73). I can not do it one by one. I want programmatic way of doing it.
A:
<code>
import numpy as np
a = np.array([[1,5,9,13,17],
              [2,6,10,14,18],
              [3,7,11,15,19],
              [4,8,12,16,20]])
patch_size = 2
</code>
result = ... # put solution in this variable
BEGIN SOLUTION
<code>

x = a[:a.shape[0] // patch_size * patch_size, :a.shape[1] // patch_size * patch_size]
result = x.reshape(x.shape[0]//patch_size, patch_size, x.shape[1]// patch_size, patch_size).swapaxes(1, 2).transpose(1, 0, 2, 3).reshape(-1, patch_size, patch_size)

import numpy as np
import copy
import tokenize, io

def generate_test_case(test_case_id):
    def define_test_input(test_case_id):
        if test_case_id == 1:
            a = np.array(
                [
                    [1, 5, 9, 13, 17],
                    [2, 6, 10, 14, 18],
                    [3, 7, 11, 15, 19],
                    [4, 8, 12, 16, 20],
                ]
            )
            patch_size = 2
        elif test_case_id == 2:
            np.random.seed(42)
            a = np.random.rand(100, 200)
            patch_size = np.random.randint(4, 8)
        return a, patch_size

    def generate_ans(data):
        _a = data
        a, patch_size = _a
        x = a[
            : a.shape[0] // patch_size * patch_size,
            : a.shape[1] // patch_size * patch_size,
        ]
        result = (
            x.reshape(
                x.shape[0] // patch_size,
                patch_size,
                x.shape[1] // patch_size,
                patch_size,
            )
            .swapaxes(1, 2)
            .transpose(1, 0, 2, 3)
            .reshape(-1, patch_size, patch_size)
        )
        return result

    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):
    np.testing.assert_array_equal(result, ans)
    return 1

exec_context = r"""
import numpy as np
a, patch_size = 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 test_string(solution: str):
    tokens = []
    for token in tokenize.tokenize(io.BytesIO(solution.encode("utf-8")).readline):
        tokens.append(token.string)
    assert "while" not in tokens and "for" not in tokens


Prompt:
Problem:
I have an array :
a = np.array([[ 0,  1,  2,  3, 5, 6, 7, 8],
              [ 4,  5,  6,  7, 5, 3, 2, 5],
              [ 8,  9, 10, 11, 4, 5, 3, 5]])
I want to extract array by its columns in RANGE, if I want to take column in range 1 until 5, It will return
a = np.array([[ 1,  2,  3, 5, ],
              [ 5,  6,  7, 5, ],
              [ 9, 10, 11, 4, ]])
How to solve it? Thanks
A:
<code>
import numpy as np
a = np.array([[ 0,  1,  2,  3, 5, 6, 7, 8],
              [ 4,  5,  6,  7, 5, 3, 2, 5],
              [ 8,  9, 10, 11, 4, 5, 3, 5]])
low = 1
high = 5
</code>
result = ... # put solution in this variable
BEGIN SOLUTION
<code>

import numpy as np
import pandas as pd
import copy

def generate_test_case(test_case_id):
    def define_test_input(test_case_id):
        if test_case_id == 1:
            a = np.array(
                [
                    [0, 1, 2, 3, 5, 6, 7, 8],
                    [4, 5, 6, 7, 5, 3, 2, 5],
                    [8, 9, 10, 11, 4, 5, 3, 5],
                ]
            )
            low = 1
            high = 5
        elif test_case_id == 2:
            np.random.seed(42)
            a = np.random.rand(20, 10)
            low = np.random.randint(1, 8)
            high = np.random.randint(low + 1, 10)
        return a, low, high

    def generate_ans(data):
        _a = data
        a, low, high = _a
        result = a[:, low:high]
        return result

    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):
    np.testing.assert_array_equal(result, ans)
    return 1

exec_context = r"""
import numpy as np
a, low, high = 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)