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

Question

^{灵感来源于如何在R？中创建数字序列的帖子}

问题:

我想在NumPy中制作以下序列。

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

我尝试了以下方法：

• Non-generic and hard coding using 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])
  

  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（没有Numpy）方法是：

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

这对于较小的n（~150）来说会更快，但对于较大的n来说会比@tangolin的解决方案慢。它仍然比OP的“纯python”方式快。

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

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

Answer 2

我不确定这是否是一个好主意，但我试图运行它对你的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