Соединить строки вдоль оси

Question

Предположим, у меня есть простой массив строк, например:

import numpy as np
print('numpy version:', np.__version__)

a = np.arange(25).reshape(5, 5)
stra = a.astype(np.dtype(str))

print(stra)

output:

numpy version: 1.15.2
[['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']]

Я хочу работать вдоль заданной оси, выбрать несколько элементов и соединитьсяэти строки.Сначала я попробовал это:

print(np.apply_along_axis('|'.join, 1, stra.take([2, 3], 1)))

Но более длинные результирующие строки усекаются, чтобы соответствовать кратчайшему:

['2|3' '7|8' '12|' '17|' '22|']

Конечно, я могу написать свой собственный цикл, чтобы получить желаемый вывод,но это несколько неудовлетворительно делать, когда однострочник почти работает.

def join_along_axis(array, indices, axis):        
    if array.ndim == 1:
        return np.array('|'.join(array.take(indices)))

    joined = []        
    # Move axis of interest to end and flatten others to make the loop easy.
    work_arr = np.rollaxis(array, axis, -1)
    shape = work_arr.shape
    new_shape = (np.product(work_arr.shape[:-1]), work_arr.shape[-1])
    work_arr = work_arr.reshape(new_shape)

    for arr in work_arr:
        joined.append('|'.join(arr.take(indices)))

    return np.array(joined).reshape(shape[:-1])

print(join_along_axis(stra, [2, 3], 1))

output:

['2|3' '7|8' '12|13' '17|18' '22|23']

Есть ли более простой способ сделать это, чем моя join_along_axis функция?

Обновление для ясности: Мне нужно, чтобы это было достаточно общим, чтобы работать с массивом с произвольным числом измерений и вдоль любой выбранной оси.

Answer 1

Из разговора GitHub, связанного с @theshopen, кажется, я могу использовать lambda, чтобы указать желаемый размер строки.Так что это работает:

lens = np.vectorize(len)
indices = [2, 3]
axis = 1

new_len = lens(stra.take(indices, axis)).sum(1).max() + len(indices) - 1
new_type = '{}{}'.format(stra.dtype.char, new_len)

print(np.apply_along_axis(
    lambda x: np.array('|'.join(x), new_type),
    axis, stra.take(indices, axis)))

Answer 2

Сначала я попытался сделать это по-своему, используя apply_along_axis, но обнаружил, что это может быть сложнее, , по-видимому, NP недостаточно хорошо подходит для работы со строками.

Так как насчет понимания списка?

a =a = np.arange(25).reshape(5, 5)
stra = a.astype(np.dtype(str))
only23 = zip(stra[:,2],stra[:,3])
only23

output:

[('2', '3'), ('7', '8'), (' 12 ',' 13 '), (' 17 ',' 18 '), (' 22 ',' 23 ')]

теперь позволяет выполнить понимание:

[x[0] +'|'+x[1] for x in only23]

вывод:

['2 | 3', '7 | 8', '12 | 13 ', '17 | 18', '22 | 23 ']

И вы действительно можете сделать его однострочным, я просто не думаю, что он будет таким читабельным