從給定包含開始和結束索引的 ndarray 的二維數組中提取

Question

data (shape: [ chn , dim ]) 是一個二維數組，其中有chn通道和dim維度。

indices (shape: [ seg ,2]) 由seg數組成，其中每個段包含一個開始和結束索引。 例如，我們要提取 4 個段，其中第一個段從3開始，到5結束（含）：

indices = np.array([[3,5], [8,10], [13,15], [16,18]])

extracted的 output（形狀：[ seg , chn , dim ]）是一個 3 維數組，其中它由seg數（根據indices ）、 chn通道和dim維度（根據data ）組成。

問題：如何定義do_something() function 以創建extracted的 numpy 數組？

data1 = np.expand_dims(np.arange(0, 100), 0)
data2 = np.expand_dims(np.arange(100, 200), 0)
data3 = np.expand_dims(np.arange(200, 300), 0)
data = np.concatenate([data1, data2, data3])
print("data.shape", data.shape) # data.shape (3, 100)
print(data)
# [[  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  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 133 134 135
#   136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153
#   154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171
#   172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189
#   190 191 192 193 194 195 196 197 198 199]
#  [200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217
#   218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235
#   236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253
#   254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271
#   272 273 274 275 276 277 278 279 280 281 282 283 284 285 286 287 288 289
#   290 291 292 293 294 295 296 297 298 299]]

indices = np.array([[3,5], [8,10], [13,15], [16,18]])

# extracted = do_something() # TODO: need help here!

extracted = np.array(
[
    [[3,4,5],[103,104,105],[203,204,205]],
    [[8,9,10],[108,109,110],[208,209,210]],
    [[13,14,15],[113,114,115],[213,214,215]],
    [[16,17,18],[116,117,118],[216,217,218]],
])
print("extracted.shape", extracted.shape) # extracted.shape (4, 3, 3)
print(extracted)
# [[[  3   4   5]
#   [103 104 105]
#   [203 204 205]]

#  [[  8   9  10]
#   [108 109 110]
#   [208 209 210]]

#  [[ 13  14  15]
#   [113 114 115]
#   [213 214 215]]

#  [[ 16  17  18]
#   [116 117 118]
#   [216 217 218]]]

Answer 1

import numpy as np

data = np.arange(300).reshape((3, 100))
indices = np.asarray([[3, 5], [8, 10], [13, 15], [16, 18]])

extracted_2d = np.concatenate([data[:, start:end + 1] for (start, end) in indices])
print(extracted_2d.shape) # (12, 3)
# can just reshape to (4, 3, 3) though
# extracted_2d.reshape((4, 3, 3))

# introduce a new dimension (1, 3, 3) by indexing with None and then concatenate those 4 results
extracted = np.concatenate([data[None, :, start:end + 1] for (start, end) in indices])
print(extracted.shape) # (4, 3, 3)

Answer 2

您可以使用：

indices = np.array([[3, 5], [8, 10], [13, 15], [16, 18]])
idx = [slice(start, stop + 1) for start, stop in indices]

extracted = np.zeros((4, 3, 3))
for i, id_ in enumerate(idx):
    extracted[i, :] = data[:, id_]

這給了你：

extracted.shape # (4, 3, 3) as you defined

並extracted為：

extracted = 
array([[[  3.,   4.,   5.],
        [103., 104., 105.],
        [203., 204., 205.]],

       [[  8.,   9.,  10.],
        [108., 109., 110.],
        [208., 209., 210.]],

       [[ 13.,  14.,  15.],
        [113., 114., 115.],
        [213., 214., 215.]],

       [[ 16.,  17.,  18.],
        [116., 117., 118.],
        [216., 217., 218.]]])