python-practice-problems/problems/numpy/select_row_col.py at main · abrahmasandra/python-practice-problems · GitHub

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
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
import numpy as np
def select_row_col(x, row_idx=None, col_idx=None):
    """
    Select a subset of rows or columns in the two-dimensional array x.
    Inputs:
        x: input two-dimensional array
        row_idx: a list of row index we are selecting, None if not specified
        col_idx: a list of column index we are selecting, None if not specified
    Returns: a two-dimensional array where we have selected based on the
        specified row_idx and col_idx
    """

    # YOUR CODE HERE
    # Replace None with an appropriate return value
    tmp = np.copy(x)
    if row_idx:
        tmp = tmp[row_idx, :]
    if col_idx:
        tmp = tmp[:, col_idx]
    return tmp

if __name__ == "__main__":
    x = np.array([[0, 1, 2],
                    [3, 4, 5],
                    [6, 7, 8]])
    print(select_row_col(x, [1, 2], None))
    print(select_row_col(x, None, [1, 2]))
    print(select_row_col(x, [1, 2], [0, 2]))


#############################################################
###                                                       ###
###                    Testing code.                      ###
###    !!! DO NOT MODIFY ANY CODE BELOW THIS POINT !!!    ###
###                                                       ###
#############################################################

import sys
import numpy as np
sys.path.append('../')

import test_utils as utils

def test_select_row_col():
    x = np.arange(20).reshape(4, 5)

    def custom_get_cols(x, cols):
        out = np.stack([x[:, r] for r in cols], -1)
        return out

    def custom_get_rows(x, rows):
        out = np.stack([x[r] for r in rows])
        return out

    def custom_get_rows_cols(x, rows, cols):
        row_out = np.stack([x[r] for r in rows])
        out = np.stack([row_out[:, r] for r in cols], -1)
        return out


    for tgt_cols in [[0], [1, 2, 3], [3, 2, 1], [2, 1, 3, 4, 0]]:

        recreate_msg = utils.gen_recreate_msg('select_row_col', x, None, tgt_cols)

        result = select_row_col(x, None, tgt_cols)

        utils.check_none(result, recreate_msg)
        utils.check_is_ndarray(result, recreate_msg)

        expected_shape = (4, len(tgt_cols))
        assert result.shape == expected_shape, \
            "The shape of the returned array was {}, but".format(result.shape) \
            + " we expected {}\n\n".format(expected_shape) + recreate_msg

        expected_value = custom_get_cols(x, tgt_cols)

        utils.check_array_equal(result, expected_value,
                          recreate_msg)

    for tgt_rows in [[0], [1, 2, 3], [3, 2, 1], [2, 1, 3, 0]]:

        recreate_msg = utils.gen_recreate_msg('select_row_col', x, tgt_rows, None)

        result = select_row_col(x, tgt_rows, None)

        utils.check_none(result, recreate_msg)
        utils.check_is_ndarray(result, recreate_msg)

        expected_shape = (len(tgt_rows), 5)
        assert result.shape == expected_shape, \
            "The shape of the returned array was {}, but".format(result.shape) \
            + " we expected {}\n\n".format(expected_shape) + recreate_msg

        expected_value = custom_get_rows(x, tgt_rows)

        utils.check_array_equal(result, expected_value,
                          recreate_msg)

    for tgt_rows, tgt_cols in [([0],[0]), ([1, 2, 3],[1, 3]), ([3, 1],[0,2])]:

        recreate_msg = utils.gen_recreate_msg('select_row_col', x, tgt_rows, tgt_cols)

        result = select_row_col(x, tgt_rows, tgt_cols)

        utils.check_none(result, recreate_msg)
        utils.check_is_ndarray(result, recreate_msg)

        expected_shape = (len(tgt_rows), len(tgt_cols))
        assert result.shape == expected_shape, \
            "The shape of the returned array was {}, but".format(result.shape) \
            + " we expected {}\n\n".format(expected_shape) + recreate_msg

        expected_value = custom_get_rows_cols(x, tgt_rows, tgt_cols)

        utils.check_array_equal(result, expected_value,
                          recreate_msg)

    for tgt_rows, tgt_cols in [(None, None)]:

        recreate_msg = utils.gen_recreate_msg('select_row_col', x, None, None)

        result = select_row_col(x, None, None)

        utils.check_none(result, recreate_msg)
        utils.check_is_ndarray(result, recreate_msg)

        expected_shape = (x.shape[0], x.shape[1])
        assert result.shape == expected_shape, \
            "The shape of the returned array was {}, but".format(result.shape) \
            + " we expected {}\n\n".format(expected_shape) + recreate_msg

        utils.check_array_equal(result, x, recreate_msg)