如何在 NumPy 中创建一系列数字序列？

Question

^{受How to create a sequence of sequences of numbers in R? 帖子的启发。}

---

问题：

我想在 NumPy 中创建以下序列。

[1, 2, 3, 4, 5, 2, 3, 4, 5, 3, 4, 5, 4, 5, 5]

我尝试了以下方法：

• 使用np.r_ 的非泛型和硬编码

  np.r_[1:6, 2:6, 3:6, 4:6, 5:6]
  # array([1, 2, 3, 4, 5, 2, 3, 4, 5, 3, 4, 5, 4, 5, 5])
  

• 纯 Python 生成所需的数组。

  n = 5
  a = np.r_[1:n+1]
  [i for idx in range(a.shape[0]) for i in a[idx:]]
  # [1, 2, 3, 4, 5, 2, 3, 4, 5, 3, 4, 5, 4, 5, 5]
  

• 创建一个二维数组并从中取出上面的三角形。

  n = 5
  a = np.r_[1:n+1]
  arr = np.tile(a, (n, 1))
  print(arr)
  # [[1 2 3 4 5]
  #  [1 2 3 4 5]
  #  [1 2 3 4 5]
  #  [1 2 3 4 5]
  #  [1 2 3 4 5]]
  
  o = np.triu(arr).flatten()
  # array([1, 2, 3, 4, 5, 
  #        0, 2, 3, 4, 5, 
  #        0, 0, 3, 4, 5, # This is 1D array
  #        0, 0, 0, 4, 5, 
  #        0, 0, 0, 0, 5])
  
  out = o[o > 0]
  # array([1, 2, 3, 4, 5, 2, 3, 4, 5, 3, 4, 5, 4, 5, 5])
  

上述解决方案是通用的，但我想知道在 NumPy 中是否有更有效的方法。

Answer 1

我不确定这是否是个好主意，但我尝试针对你的 python 方法运行它，它似乎更快。

np.concatenate([np.arange(i, n+1) for i in range(1, n+1)])

这里是完整的代码：

import numpy as np
from time import time

n = 5000

t = time()
c = np.concatenate([np.arange(i, n+1) for i in range(1, n+1)])
print(time() - t)
# 0.039876699447631836

t = time()
a = np.r_[1:n+1]
b = np.array([i for idx in range(a.shape[0]) for i in a[idx:]])
print(time() - t)
# 2.0875167846679688

print(all(b == c))
# True

Answer 2

一个真的纯Python（没有numpy）方式是：

n = 5
a = [r for start in range(1, n+1) for r in range(start, n+1)]

对于较小的 n (~150)，这会更快，但对于较大的 n，这将比 @tangolin 的 solution 慢。它仍然比OP的“纯python”方式更快。

更快的实现会提前准备数据，避免每次都创建新范围：

source = np.arange(1, n+1)
d = np.concatenate([source[i: n+1] for i in range(0, n)])

注意

我原来的实现既为返回值分配空间，又提前准备好数据，但不是pythonic。在阅读@tangolin's answer 后，我将其更改为使用concatenate，并注意到concatenate 也是如此。

原始实现：

e = np.empty((n*(n+1)//2, ), dtype='int64')
source = np.arange(1, n+1)
for i in range(n):
    init = n * i - i*(i-1)//2
    end = n - i + init
    e[init:end] = source[i:n]