index() will give the first occurrence of an item in a list. Is there a neat trick which returns all indices in a list for an element?
17 Answers
You can use a list comprehension:
indices = [i for i, x in enumerate(my_list) if x == "whatever"]
The iterator enumerate(my_list) yields pairs (index, item) for each item in the list. Using i, x as loop variable target unpacks these pairs into the index i and the list item x. We filter down to all x that match our criterion, and select the indices i of these elements.
While not a solution for lists directly, numpy really shines for this sort of thing:
import numpy as np
values = np.array([1,2,3,1,2,4,5,6,3,2,1])
searchval = 3
ii = np.where(values == searchval)[0]
returns:
ii ==>array([2, 8])
This can be significantly faster for lists (arrays) with a large number of elements vs some of the other solutions.
A solution using list.index:
def indices(lst, element):
result = []
offset = -1
while True:
try:
offset = lst.index(element, offset+1)
except ValueError:
return result
result.append(offset)
It's much faster than the list comprehension with enumerate, for large lists. It is also much slower than the numpy solution if you already have the array, otherwise the cost of converting outweighs the speed gain (tested on integer lists with 100, 1000 and 10000 elements).
NOTE: A note of caution based on Chris_Rands' comment: this solution is faster than the list comprehension if the results are sufficiently sparse, but if the list has many instances of the element that is being searched (more than ~15% of the list, on a test with a list of 1000 integers), the list comprehension is faster.
more_itertools.locate finds indices for all items that satisfy a condition.
from more_itertools import locate
list(locate([0, 1, 1, 0, 1, 0, 0]))
# [1, 2, 4]
list(locate(['a', 'b', 'c', 'b'], lambda x: x == 'b'))
# [1, 3]
more_itertools is a third-party library > pip install more_itertools.
Getting all the occurrences and the position of one or more (identical) items in a list
With enumerate(alist) you can store the first element (n) that is the index of the list when the element x is equal to what you look for.
>>> alist = ['foo', 'spam', 'egg', 'foo']
>>> foo_indexes = [n for n,x in enumerate(alist) if x=='foo']
>>> foo_indexes
[0, 3]
>>>
Let's make our function findindex
This function takes the item and the list as arguments and return the position of the item in the list, like we saw before.
def indexlist(item2find, list_or_string):
"Returns all indexes of an item in a list or a string"
return [n for n,item in enumerate(list_or_string) if item==item2find]
print(indexlist("1", "010101010"))
Output
[1, 3, 5, 7]
Simple
for n, i in enumerate([1, 2, 3, 4, 1]):
if i == 1:
print(n)
Output:
0
4
- There’s an answer using
np.whereto find the indices of a single value - Converting the
listto anarray, and usingnp.whereis6.8xfaster than any list-comprehension for finding all indices of a single element. - Faster solutions using
numpycan be found in Get a list of all indices of repeated elements in a numpy array
import numpy as np
import random # to create test list
# create sample list
random.seed(365)
l = [random.choice(['s1', 's2', 's3', 's4']) for _ in range(20)]
# convert the list to an array for use with these numpy methods
a = np.array(l)
# create a dict of each unique entry and the associated indices
idx = {v: np.where(a == v)[0].tolist() for v in np.unique(a)}
# print(idx)
{'s1': [7, 9, 10, 11, 17],
's2': [1, 3, 6, 8, 14, 18, 19],
's3': [0, 2, 13, 16],
's4': [4, 5, 12, 15]}
# find a single element with
idx = np.where(a == 's1')
print(idx)
[out]:
(array([ 7, 9, 10, 11, 17], dtype=int64),)
%timeit
# create 2M element list
random.seed(365)
l = [random.choice(['s1', 's2', 's3', 's4']) for _ in range(2000000)]
# create array
a = np.array(l)
Find the indices of one value
- Find indices of a single element in a 2M element list with 4 unique elements
# np.where
%timeit np.where(a == 's1')
[out]:
25.9 ms ± 827 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)
# list-comprehension
%timeit [i for i, x in enumerate(l) if x == "s1"]
[out]:
175 ms ± 2.73 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)
# filter
%timeit list(filter(lambda i: a[i]=="s1", range(len(a))))
[out]:
925 ms ± 10.7 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
Find the indices of all the values
- Find indices of all unique elements in a 2M element list with 4 unique elements
# use np.where and np.unique
%timeit {v: np.where(a == v)[0].tolist() for v in np.unique(a)}
[out]:
312 ms ± 9.13 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
# list comprehension inside dict comprehension
%timeit {req_word: [idx for idx, word in enumerate(a) if word == req_word] for req_word in set(a)}
[out]:
3.72 s ± 21.7 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
Using a for-loop:
- Answers with
enumerateand a list comprehension are more pythonic, not necessarily faster. However, this answer is aimed at students who may not be allowed to use some of those built-in functions. - create an empty list,
indices - create the loop with
for i in range(len(x)):, which essentially iterates through a list of index locations[0, 1, 2, 3, ..., len(x)-1] - in the loop, add any
i, wherex[i]is a match tovalue, toindices
def get_indices(x: list, value: int) -> list:
indices = list()
for i in range(len(x)):
if x[i] == value:
indices.append(i)
return indices
n = [1, 2, 3, -50, -60, 0, 6, 9, -60, -60]
print(get_indices(n, -60))
>>> [4, 8, 9]
- The functions,
get_indices, are implemented with type hints. In this case, the list,n, is a bunch ofints, therefore we search forvalue, also defined as anint.
Using a while-loop and .index:
- With
.index, usetry-exceptfor error handling, because aValueErrorwill occur ifvalueis not in thelist.
def get_indices(x: list, value: int) -> list:
indices = list()
i = 0
while True:
try:
# find an occurrence of value and update i to that index
i = x.index(value, i)
# add i to the list
indices.append(i)
# advance i by 1
i += 1
except ValueError as e:
break
return indices
print(get_indices(n, -60))
>>> [4, 8, 9]
A dynamic list comprehension based solution incase we do not know in advance which element:
lst = ['to', 'be', 'or', 'not', 'to', 'be']
{req_word: [idx for idx, word in enumerate(lst) if word == req_word] for req_word in set(lst)}
results in:
{'be': [1, 5], 'or': [2], 'to': [0, 4], 'not': [3]}
You can think of all other ways along the same lines as well but with index() you can find only one index although you can set occurrence number yourself.
Here is a time performance comparison between using np.where vs list_comprehension. Seems like np.where is faster on average.
# np.where
start_times = []
end_times = []
for i in range(10000):
start = time.time()
start_times.append(start)
temp_list = np.array([1,2,3,3,5])
ixs = np.where(temp_list==3)[0].tolist()
end = time.time()
end_times.append(end)
print("Took on average {} seconds".format(
np.mean(end_times)-np.mean(start_times)))
Took on average 3.81469726562e-06 seconds
# list_comprehension
start_times = []
end_times = []
for i in range(10000):
start = time.time()
start_times.append(start)
temp_list = np.array([1,2,3,3,5])
ixs = [i for i in range(len(temp_list)) if temp_list[i]==3]
end = time.time()
end_times.append(end)
print("Took on average {} seconds".format(
np.mean(end_times)-np.mean(start_times)))
Took on average 4.05311584473e-06 seconds