Function for setting a streak of strings in a Pandas I have a df with artist names in one column, I want to find how many times the artists repeats through the column.

Example:

df = pd.DataFrame({"Artist":["Tim Maia","Jorge Ben","Tim Maia",
                             "Tim Maia","Jorge Ben","Roberto Carlos",
                             "Roberto Carlos","Roberto Carlos",
                             "Roberto Carlos",
                             "Dire Straits"]})


I've already had some success using itertuples(), so that was my starting point for this function:

def artist_streak(df):
    # created to check if artists repeats, not returned
    artists = []
    # this will be returned
    streaks = []
    # since each observation counts, makes sense to start at 1
    count=1

    for row in df.itertuples():
        # artists[-1] in an empty list will raise IndexError:
        try:
            # Checking if current artist is the same as the last appended
            if artists[-1] == row.Artist:
                    count += 1
                    artists.append(row.Artist)
                    streaks.append(count)
            else:
                # The artist does not repeated  
                count = 1
                artists.append(row.Artist)
                streaks.append(count)
        # Here we set things right for the first iteration
        except IndexError:
            artists.append(row.Artist)
            streaks.append(count)


    return streaks


Usage:

df["Result"] = artist_streak(df)


My concerns about this code are:

1. The artists list [ ]:

It is the way I could think of checking the current iterator item vs. the previous one, I don't know if it is a good way.

2. The try/except part:

As it addresses the specific case of list artists being empty, I don't know if it hurts performance (i.e. there is a faster way to do it).

Any other suggestions on how to make it faster, cleaner, or even pythonic will be welcome.
For starter, your implementation is too tightly tied to your data, making it not reusable for other kind of dataframe where you would want to count streaks. If you provided a column (pd.Series) to the function instead, it would be much more versatile; allowing calls such as count_streak(df.Artist), count_streak(pd.Series([1, 2, 2, 1, 1, 3]) or even count_streak('abcdccbdddbbbbcaaacbddba').

Doing so, you would remove the need for itertuple and favor regular iteration. This would ease the use of enumerate to check for the previous element:

def count_streak(sequence):
    streaks = []

    for index, element in enumerate(sequence):
        if not index or sequence[index - 1] != element:
            count = 1
        else:
            count += 1
        streaks.append(count)

    return streaks


Much simpler. But this is all without taking into account the wonderful itertools module and its groupby function. It will make group of consecutive data that evaluate equals, yielding both the value, and the entire group. The other advantage being that any iterable will be accepted:

def count_streak(iterable):
    return [
        i for _, g in itertools.groupby(iterable)
        for i, _ in enumerate(g, start=1)
    ]


Since you wish to stay in pandas' world, you also have a groupby method available on DataFrames and Series. But contrary to itertools.groupby that stops at each boundaries, pandas' groupby return all occurences of similar data even if some are discontinuous. But you can make use of cumsum and cumcount to get an equivalent and vectorize these operation instead of a slow python loop. This would require a pandas.Series as parameter, though, not any iterable:

def count_streak(serie):
    return serie.groupby((serie != serie.shift()).cumsum()).cumcount() + 1


usage being

df['Streak'] = count_streak(df.Artist)




And since you commented on performances, here is a small benchmark comparing the approaches. I’ve thrown in a couple extra approaches, out of curiosity. The code to produce the results is:

import pandas as pd
import numpy as np
import timeit

def python_streaks(sequence):
  streaks = []
  for index, element in enumerate(sequence):
    if not index or sequence[index - 1] != element:
      count = 1
    else:
      count += 1
    streaks.append(count)
  return streaks

def python2_streaks(sequence):
  streaks = [1] * len(sequence)
  for index, element in enumerate(sequence):
    if index and sequence[index - 1] == element:
      streaks[index] += streaks[index - 1]
  return streaks

def numpy_streaks(sequence):
  streaks = pd.Series(np.ones(len(sequence), dtype=int))
  for index, element in enumerate(sequence):
    if index and sequence[index - 1] == element:
      streaks[index] += streaks[index - 1]
  return streaks

def itertools_streak(iterable):
  return [
    i for _, g in itertools.groupby(iterable)
    for i, _ in enumerate(g, start=1)
  ]

def pandas_streaks(serie):
  return serie.groupby((serie != serie.shift()).cumsum()).cumcount() + 1

for size in range(8):
  size = 10**size
  for name in ('python', 'python2', 'numpy', 'pandas', 'itertools'):
    setup = 'import pandas as pd;from random import choice;from __main__ import {}_streaks as streaks; df = pd.DataFrame({{"Artist": [choice("abcde") for _ in range({})]}})'.format(name, size)
    elapsed = timeit.timeit('df["Streaks"] = streaks(df.Artist)', setup=setup, number=1000)
    print(name, size, elapsed)


And the results are (I stopped pure Python loops as they were taking too much time):



or, without using log scale on the Y-axis:



So depending on the amount of data you wish to process, either itertools (up to 10000ish elements) or pandas version is better.