How to count adjacent recurring elements in an array?

Viewed 616

I have an array of 0's and 1's as such

[0,0,1,1,1,0,0,0,0,1,1,0,0]

I want to define a function that will take this array as input and output an array of the same length, with the count of adjacent 1's in the index where the first 1 appeared (and 0 otherwise). So the output would be

[0,0,3,0,0,0,0,0,0,2,0,0,0]

because 1 appeared in the 2nd index 3 consecutive times and 1 appeared in the 9th index 2 consecutive times.

Is there a way to do this using numpy? If not, is there some (efficient) pythonic way to do it?

8 Answers

Here is a solution using pure vectorized operations and no list iterations:

import numpy as np

data = np.array([0,0,1,1,1,0,0,0,0,1,1,0,0])
output = np.zeros_like(data)

where = np.where(np.diff(data))[0]
vals = where[1::2] - where[::2]
idx = where[::2] + 1

output[idx] = vals
output
# array([0, 0, 3, 0, 0, 0, 0, 0, 0, 2, 0, 0, 0])

Using the itertools module:

from itertools import chain, groupby

A = [0,0,1,1,1,0,0,0,0,1,1,0,0]

def get_lst(x):
    values = list(x[1])
    return [len(values)] + [0]*(len(values) - 1) if x[0] else values

res = list(chain.from_iterable(map(get_lst, groupby(A))))

# [0, 0, 3, 0, 0, 0, 0, 0, 0, 2, 0, 0, 0]

You could use groupby to group the consecutive elements:

from itertools import groupby

a = [0, 0, 1, 1, 1, 0, 0, 0, 0, 1, 1, 0, 0]


def groups(lst):
    result = []
    for key, group in groupby(lst):
        if not key:  # is a group of zeroes
            result.extend(list(group))
        else:  # is a group of ones
            count = sum(1 for _ in group)
            if count > 1:  # if more than one
                result.append(count)
                result.extend(0 for _ in range(count - 1))
            else:
                result.append(0)
    return result


print(groups(a))

Output

[0, 0, 3, 0, 0, 0, 0, 0, 0, 2, 0, 0, 0]

A shorter (more pythonic?) is the following:

def groups(lst):
    for key, group in groupby(lst):
        count = sum(1 for _ in group)
        if key and count > 1:
            yield count
        yield from (0 for _ in range(count - key))


print(list(groups(a)))

Here's one way using numpy and a list comprehension:

In [23]: a = np.array([0,0,1,1,1,0,0,0,0,1,1,0,0])
In [24]: np.hstack([x.sum() if x[0] == 1 else x for x in np.split(a, np.where(np.diff(a) != 0)[0]+1)])
Out[24]: array([0, 0, 3, 0, 0, 0, 0, 2, 0, 0])

The logic:

  1. Find leading and trailing indices of where you have consequence 1s.
  2. Split your array from those indices
  3. sum those sub lists that have one and leave sub lists with zero the way they are
  4. flatten the result using np.hstack.

If you want to replace the remained ones with 0 just do the following:

In [28]: np.hstack([[x.sum(), *[0]*(len(x) -1)]  if x[0] == 1 else x for x in np.split(a, np.where(np.diff(a) != 0)[0]+1)])
Out[28]: array([0, 0, 3, 0, 0, 0, 0, 0, 0, 2, 0, 0, 0])

[0]*(len(x) -1) will create the expected 0s for you and using an in-place unpacking you'll be able to place them beside the sum(x).

If you ever wanted a pure Python approach here's one way using itertools.groupby:

In [63]: def summ_cons(li):
    ...:     for k,g in groupby(li) :
    ...:            if k:
    ...:               s = sum(g)
    ...:               yield s
    ...:               yield from (0 for _ in range(s-1))
    ...:            yield from g
    ...:            


In [65]: list(summ_cons(a))
Out[65]: [0, 0, 3, 0, 0, 0, 0, 0, 0, 2, 0, 0, 0]

Use pandas, take advantage in pandas count counts non-NaN values. Create NaNs using mask then groupby changes in values of s.

import pandas as pd
l = [0,0,1,1,1,0,0,0,0,1,1,0,0]
s = pd.Series(l)
g = s.diff().ne(0).cumsum()
s.mask(s==0).groupby(g).transform('count').mask(g.duplicated(), 0).tolist()

Output:

[0, 0, 3, 0, 0, 0, 0, 0, 0, 2, 0, 0, 0]

using groupby

from itertools import groupby
a=[0,0,1,1,1,0,0,0,0,1,1,0,0]
lst_of_tuples=[]
for k,v in groupby(a):
    lst_of_tuples.append((k,len(list(v))))
ans=[]
for k,v in lst_of_tuples:
    temp=[v if k==1 else k]
    for i in range(v-1):
        temp.append(0)
    ans=ans+temp

output

[0,0,3,0,0,0,0,0,0,2,0,0,0]

TL;DR: this gives you the output you want:

import itertools

input = [0,0,1,1,1,0,0,0,0,1,1,0,0]   
result = []

for k, g in itertools.groupby(input):
    if k == 1:
        ll = len(list(g))
        result.extend([ll,] + [0 for _ in range(ll-1)])
    else:
        result.extend(list(g)) 

will give you:

[0, 0, 3, 0, 0, 0, 0, 0, 0, 2, 0, 0, 0]

explanation:

itertools has the groupby function to split runs of "the same"

for k, g in itertools.groupby(input):
    print(g, list(k))

will give you:

0 [0, 0]
1 [1, 1, 1]
0 [0, 0, 0, 0]
1 [1, 1]
0 [0, 0]

so k is the key, the element in the input sequence and g is the group.

so the if conditions appends either (in case of a 0) the series of 0's from the input, or the length if the run of 1's plus zeros to fill up for the length of the 1-run.

Other option with no dependencies:

The good old while loop accessing by index (sometimes faster than numpy):

def count_same_adjacent_non_zeros(iterable):
  i, x, size = 0, 0, len(iterable)
  while i < size-1:
    if iterable[i] != iterable[i+1]:
      tmp = iterable[x:i+1]
      if not iterable[i] == 0:
        tmp = [len(tmp)] + [0 for _ in range(i-x)]
      for e in tmp: yield e
      x = i + 1
    i += 1
  for e in iterable[x:size]: yield e


print(list(count_same_adjacent_non_zeros(array)))

#=> [0, 0, 3, 0, 0, 0, 0, 0, 0, 2, 0, 0, 0]

It works also with array = [0,0,4,4,4,0,0,0,0,5,5,0,0]

Related