Comment obtenir des indices de concaténation numpy

Question

Je veux concatenate numpy ndarray s mais avec des indices aussi.

C'est-à-dire pour

>>> np.concatenate([[1,2,3,4], [11,22,33], [4], [435,456]]) 

Je veux que le résultat numpy.concatenate donne déjà:

array([ 1, 2, 3, 4, 11, 22, 33, 4, 435, 456]) 

mais je aussi veux que les indices de la liste d'entrée indcating quel élément de l'entrée liste chaque élément du tableau de sortie provient de:

array([0, 0, 0, 0, 1, 1, 1, 2, 3, 3] 

Je connais des recettes simples basées sur np.cumsum et np.repeat que je pourrais utiliser pour construire le tableau d'index moi-même, mais y a-t-il une manière numpy-native de le faire?

Answer 1

est ici une solution basée sur la liste simple:

In [64]: alist = [[1,2,3,4], [11,22,33], [4], [435,456]] 
In [67]: idlist = [[i+1]*len(x) for i,x in enumerate(alist)] 
In [68]: np.concatenate(alist) 
Out[68]: array([ 1, 2, 3, 4, 11, 22, 33, 4, 435, 456]) 
In [69]: np.concatenate(idlist) 
Out[69]: array([1, 1, 1, 1, 2, 2, 2, 3, 4, 4]) 

Si les sous-listes étaient vraiment long, ou étaient déjà des réseaux, une solution en utilisant plus numpy opérations pourrait être plus rapide. Mais si c'est le mélange typique de sous-listes, la compréhension de la liste sera probablement plus rapide. Rappelez-vous que l'entrée concatenate est une liste (même si l'entrée est un tableau, il est traité comme une liste.)

In [73]: timeit idlist = [np.arange(len(x)) for i,x in enumerate(alist)] 
6.06 µs ± 224 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each) 
In [74]: timeit idlist = [[i+1]*len(x) for i,x in enumerate(alist)] 
2.98 µs ± 8.3 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each) 

Answer 2

Comment cela?

In [16]: alist = [[1,2,3,4], [11,22,33], [4], [435,456]] 

In [17]: my_arrays = [np.array(x) for x in alist] 

In [18]: index_arrays = [np.ones((x.shape), dtype=int)*i for i, x in enumerate(my_arrays)] 

In [19]: np.concatenate(my_arrays) 
Out[19]: array([ 1, 2, 3, 4, 11, 22, 33, 4, 435, 456]) 

In [20]: np.concatenate(index_arrays) 
Out[20]: array([0, 0, 0, 0, 1, 1, 1, 2, 3, 3])