Как в векторе чисел посчитать количество положительных или отрицательных элементов в ряд?

Question

Есть вектор с числами.

import numpy as np
arr = np.array([3, 2, 1, 1, -1, -2, -3, 1, 1, 2, 1], float)

Задача посчитать количество положительных или отрицательных чисел в ряд векторизированным способом. Сбрасывать до 1 если изменился знак. Появление нуля в векторе маловероятно.

Ожидаемый результат:

[nan 2. 3. 4. 1. 2. 3. 1. 2. 3. 4.]

Пример решения с циклом:

arr = np.array([3, 2, 1, 1, -1, -2, -3, 1, 1, 2, 1], float)

counter = 1
res = np.full_like(arr, np.nan)

for i in range(len(arr)):
    if i > 0:
        if np.sign(arr[i]) == np.sign(arr[i - 1]):
            counter += 1
            res[i] = counter
        else:
            counter = 1
            res[i] = counter

print(arr)
print(res)

Answer 1

Не уверен, что это самый оптимальный способ, но векторизовать можно попробовать так:

import numpy as np

arr = np.array([3, 2, 1, 1, -1, -2, -3, 1, 1, 2, 1], float)

idx = np.argwhere(np.diff(np.sign(arr))!= 0)
res = np.concatenate([np.cumsum(x.astype(bool)) for x in np.split(arr, np.ravel(idx+1))]).astype(float)

res, соответственно:

array([1., 2., 3., 4., 1., 2., 3., 1., 2., 3., 4.])

Answer 2

Без numpy:

from itertools import groupby
a = [3, 2, 1, 1, -1, -2, -3, 1, 1, 2, 1]
b = [i for _, g in groupby(a, lambda x: x>0) for i in range(1, len(list(g))+1) ]
print(b)

БЛОГ НА HUSL

Как в векторе чисел посчитать количество положительных или отрицательных элементов в ряд?

Ответы (2 шт):