Python - Split array into multiple arrays

Question

I have a array contains file names like below:

['001_1.png', '001_2.png', '001_3.png', '002_1.png','002_2.png', '003_1.png', '003_2.png', '003_3.png', '003_4.png', ....]

I want to quickly group these files into multiple arrays like this:

[['001_1.png', '001_2.png', '001_3.png'], ['002_1.png', '002_2.png'], ['003_1.png', '003_2.png', '003_3.png', '003_4.png'], ...]

Could anyone tell me how to do it in few lines in python?

Answer 1

If your data is already sorted by the file name, you can use itertools.groupby:

files = ['001_1.png', '001_2.png', '001_3.png', '002_1.png','002_2.png',
        '003_1.png', '003_2.png', '003_3.png']

import itertools

keyfunc = lambda filename: filename[:3]

# this creates an iterator that yields `(group, filenames)` tuples,
# but `filenames` is another iterator
grouper = itertools.groupby(files, keyfunc)

# to get the result as a nested list, we iterate over the grouper to
# discard the groups and turn the `filenames` iterators into lists
result = [list(files) for _, files in grouper]

print(list(result))
# [['001_1.png', '001_2.png', '001_3.png'],
#  ['002_1.png', '002_2.png'],
#  ['003_1.png', '003_2.png', '003_3.png']]

---

Otherwise, you can base your code on this recipe, which is more efficient than sorting the list and then using groupby.

• Input: Your input is a flat list, so use a regular ol' loop to iterate over it:

  for filename in files:
  

• Group identifier: The files are grouped by the first 3 letters:

  group = filename[:3]
  

• Output: The output should be a nested list rather than a dict, which can be done with

  result = list(groupdict.values())
  

Putting it together:

files = ['001_1.png', '001_2.png', '001_3.png', '002_1.png','002_2.png',
        '003_1.png', '003_2.png', '003_3.png']

import collections

groupdict = collections.defaultdict(list)
for filename in files:
    group = filename[:3]
    groupdict[group].append(filename)

result = list(groupdict.values())

print(result)
# [['001_1.png', '001_2.png', '001_3.png'],
#  ['002_1.png', '002_2.png'],
#  ['003_1.png', '003_2.png', '003_3.png']]

Read the recipe answer for more details.

Answer 2

Something like that should work:

import itertools


mylist = [...]
[list(v) for k,v in itertools.groupby(mylist, key=lambda x: x[:3])]

If input list isn't sorted, than use something like that:

import itertools


mylist = [...]
keyfunc = lambda x:x[:3]
mylist = sorted(mylist, key=keyfunc)
[list(v) for k,v in itertools.groupby(mylist, key=keyfunc)]