基于位置向numpy数组插入元素并扩展数组的方法

问题描述 #

现有一个大小为1253的numpy数组X（示例中所有元素均为1），需要将其扩展为大小1521的新数组。要求：

• 新数组中原数组元素保留在对应位置（值为1）；
• 以下指定位置的元素设为0：

array([   0,    1,    2,    3,    4,    5,    6,    7,    8,    9,   10,
     11,   12,   26,   27,   28,   29,   30,   31,   32,   33,   34,
     35,   36,   37,   38,   39,   40,   41,   42,   43,   44,   45,
     46,   47,   48,   49,   67,   68,   69,   70,   71,   72,   73,
     74,   75,   76,   77,   78,   79,   80,   81,   82,   83,   84,
     85,   86,  108,  109,  110,  111,  112,  113,  114,  115,  116,
    117,  118,  119,  120,  121,  122,  123,  149,  150,  151,  152,
    153,  154,  155,  156,  157,  158,  159,  160,  161,  189,  190,
    191,  192,  193,  194,  195,  196,  197,  198,  199,  229,  230,
    231,  232,  233,  234,  235,  236,  237,  269,  270,  271,  272,
    273,  274,  275,  309,  310,  311,  312,  313,  314,  348,  349,
    350,  351,  352,  388,  389,  390,  391,  427,  428,  429,  467,
    468,  506, 1014, 1052, 1053, 1091, 1092, 1093, 1129, 1130, 1131,
   1132, 1168, 1169, 1170, 1171, 1172, 1206, 1207, 1208, 1209, 1210,
   1211, 1245, 1246, 1247, 1248, 1249, 1250, 1251, 1283, 1284, 1285,
   1286, 1287, 1288, 1289, 1290, 1291, 1321, 1322, 1323, 1324, 1325,
   1326, 1327, 1328, 1329, 1330, 1331, 1359, 1360, 1361, 1362, 1363,
   1364, 1365, 1366, 1367, 1368, 1369, 1370, 1371, 1397, 1398, 1399,
   1400, 1401, 1402, 1403, 1404, 1405, 1406, 1407, 1408, 1409, 1410,
   1411, 1412, 1434, 1435, 1436, 1437, 1438, 1439, 1440, 1441, 1442,
   1443, 1444, 1445, 1446, 1447, 1448, 1449, 1450, 1451, 1452, 1453,
   1471, 1472, 1473, 1474, 1475, 1476, 1477, 1478, 1479, 1480, 1481,
   1482, 1483, 1484, 1485, 1486, 1487, 1488, 1489, 1490, 1491, 1492,
   1493, 1494, 1508, 1509, 1510, 1511, 1512, 1513, 1514, 1515, 1516,
   1517, 1518, 1519, 1520], dtype=int64)

解决方案 #

可以通过numpy的集合运算快速定位原数组的目标位置，再完成赋值：

import numpy as np

# 示例原数组
X = np.ones(1253)
# 指定位置数组
mask_pos = np.array([   0,    1,    2,    3,    4,    5,    6,    7,    8,    9,   10,
     11,   12,   26,   27,   28,   29,   30,   31,   32,   33,   34,
     35,   36,   37,   38,   39,   40,   41,   42,   43,   44,   45,
     46,   47,   48,   49,   67,   68,   69,   70,   71,   72,   73,
     74,   75,   76,   77,   78,   79,   80,   81,   82,   83,   84,
     85,   86,  108,  109,  110,  111,  112,  113,  114,  115,  116,
    117,  118,  119,  120,  121,  122,  123,  149,  150,  151,  152,
    153,  154,  155,  156,  157,  158,  159,  160,  161,  189,  190,
    191,  192,  193,  194,  195,  196,  197,  198,  199,  229,  230,
    231,  232,  233,  234,  235,  236,  237,  269,  270,  271,  272,
    273,  274,  275,  309,  310,  311,  312,  313,  314,  348,  349,
    350,  351,  352,  388,  389,  390,  391,  427,  428,  429,  467,
    468,  506, 1014, 1052, 1053, 1091, 1092, 1093, 1129, 1130, 1131,
   1132, 1168, 1169, 1170, 1171, 1172, 1206, 1207, 1208, 1209, 1210,
   1211, 1245, 1246, 1247, 1248, 1249, 1250, 1251, 1283, 1284, 1285,
   1286, 1287, 1288, 1289, 1290, 1291, 1321, 1322, 1323, 1324, 1325,
   1326, 1327, 1328, 1329, 1330, 1331, 1359, 1360, 1361, 1362, 1363,
   1364, 1365, 1366, 1367, 1368, 1369, 1370, 1371, 1397, 1398, 1399,
   1400, 1401, 1402, 1403, 1404, 1405, 1406, 1407, 1408, 1409, 1410,
   1411, 1412, 1434, 1435, 1436, 1437, 1438, 1439, 1440, 1441, 1442,
   1443, 1444, 1445, 1446, 1447, 1448, 1449, 1450, 1451, 1452, 1453,
   1471, 1472, 1473, 1474, 1475, 1476, 1477, 1478, 1479, 1480, 1481,
   1482, 1483, 1484, 1485, 1486, 1487, 1488, 1489, 1490, 1491, 1492,
   1493, 1494, 1508, 1509, 1510, 1511, 1512, 1513, 1514, 1515, 1516,
   1517, 1518, 1519, 1520], dtype=np.int64)

# 初始化新数组为全0
new_arr = np.zeros(1521)
# 计算新数组中需要填充原数组的位置：排除指定位置的所有索引
target_positions = np.setdiff1d(np.arange(1521), mask_pos)
# 将原数组的值赋值到目标位置
new_arr[target_positions] = X

# 验证结果（可选）
assert np.all(new_arr[mask_pos] == 0), "指定位置未全部设为0"
assert np.all(new_arr[target_positions] == 1), "原数组元素未正确填充"

关键步骤说明 #

• np.setdiff1d：计算两个数组的差集，这里用来找出新数组中不属于指定位置的索引，这些索引的数量正好等于原数组的长度（1521-268=1253）；
• 初始化全0数组：直接让指定位置保持0，无需额外修改；
• 赋值操作：将原数组元素批量填充到目标位置，高效完成扩展。

内容的提问来源于stack exchange，提问作者Yehui He