How to add count for each unique val in list

Viewed 1408

Assume I have a list.

temp = ['A', 'B', 'A', 'B', 'A', 'B']

I am looking for a way to join the count of the string inside.

Intended Output:

['A_1', 'B_1', 'A_2', 'B_2', 'A_3', 'B_3']

I was able to solve it by using a list comprehension but I am looking for a way where I don't have to specify the list [1, 1, 2, 2, 3, 3]. Is it possible?

[j + "_" + str(i) for i, j in zip([1, 1, 2, 2, 3, 3], temp)]
4 Answers

You can use collections.defaultdict with a for loop:

from collections import defaultdict

L = ['A', 'B', 'A', 'B', 'A', 'B']

dd = defaultdict(int)

res = []
for item in L:
    dd[item] += 1
    res.append(f'{item}_{dd[item]}')

print(res)

['A_1', 'B_1', 'A_2', 'B_2', 'A_3', 'B_3']

You can use a Counter or a defaultdict(int) to keep track of how many times a character has been seen as you encounter them.

>>> from collections import Counter
>>> 
>>> temp = ['A', 'B', 'A', 'B', 'A', 'B']
>>> seen = Counter()
>>> 
>>> result = []
>>> for c in temp:
...:    seen.update(c)
...:    result.append('{}_{}'.format(c, seen[c]))
...:    
>>> result
>>> ['A_1', 'B_1', 'A_2', 'B_2', 'A_3', 'B_3']

Note that seen.update(c) might have unexpected results if you expect strings with more than one character in temp. Demo:

>>> seen = Counter()
>>> seen.update('ABC')
>>> seen
>>> Counter({'A': 1, 'B': 1, 'C': 1})

Depending on how you want to count and what kind of data you expect, you might want to use the line

seen[c] += 1

instead of

seen.update(c)

Alternatively, without any imports:

>>> seen = {}
>>> result = []
>>> 
>>> for c in temp:
...:    seen[c] = seen.get(c, 0) + 1
...:    result.append('{}_{}'.format(c, seen[c]))
...:    
>>> result
>>> ['A_1', 'B_1', 'A_2', 'B_2', 'A_3', 'B_3']

You can use a dictionary (or better yet, a collections.defaultdict) to maintain the counts for each item:

from collections import defaultdict

lst = ['A', 'B', 'A', 'B', 'A', 'B']
lst2 = []
d = defaultdict(int)

for item in lst:
    d[item] += 1
    lst2.append('{}_{}'.format(item, d[item]))

print(lst2)  # ['A_1', 'B_1', 'A_2', 'B_2', 'A_3', 'B_3']

To use a list comprehension here you'd need some way to update state (i.e. the counters) for each item as you iterate the original list. For that you could use a function with a default argument, e.g.:

def get_count(item, d=defaultdict(int)):
    d[item] += 1
    return '{}_{}'.format(item, d[item])

lst2 = [get_count(item) for item in lst]
print(lst2)  # ['A_1', 'B_1', 'A_2', 'B_2', 'A_3', 'B_3']

It is in fact possible (as OP has asked in some of the comments) to do this with just a list comprehension, without too many undesirable side effects. I'm not sure whether this is necessarily a good idea though - some people may not find it the easiest code to understand:

from collections import defaultdict
import itertools

temp = ['A', 'B', 'A', 'B', 'A', 'B']

result = [j + "_" + str(next(c[j]))
          for c in [defaultdict(itertools.count)]
          for j in temp]
Related