CODEWITHSUNDEEP
pythonsetiterableenumerate
scotty3785
scotty3785

Reputation: 7006

Python: Split list in to sets

I'm trying to split a list in to enumerated sets and I have a implementation that seems to do what I want but it doesn't feel very "pythonic" is there a better way?

The below example splits the item in to sets of maximum size 3 and returns each item from items along with a set number and set index

items = [1,2,3,4,5,6,7,8,9]

def SplitSets(iterable,set_size):
    row = 0
    col = 0
    for elem in iterable:
        yield row, col, elem
        if col == set_size - 1:
            row = row+1
        col = (col + 1) % set_size

for a,b,c in SplitSets(items,3):
    print a,b,c

Expected output is

1,2,3
4,5,6
7,8,9

Upvotes: 0

Views: 587

Answers (2)

Hai Vu
Hai Vu

Reputation: 40733

One solution is to use itertools:

from itertools import count, cycle, izip

def split_sets(iterable, set_size):
    for a, b, c in izip(count(), cycle(range(set_size)), items):
        yield a // set_size, b, c

items = [1, 2, 3, 4, 5, 6, 7, 8, 9]
set_size = 3
for a, b, c in split_sets(items, set_size):
    print a, b, c

Discussion

  • For the first column (a), we call on count(), which returns 0, 1, 2, 3, 4... We then integer divide a with the set_size to get what we want
  • For the second column, range(set_size) returns [0, 1, 2] for set_size of 3. The cycle() function then repeat this sequence over and over.
  • No explanation needed for the last column
  • Now, if we take these columns and izip (which is more efficent than zip for Python 2.x) together, we will get what we want.

Upvotes: 1

Jose Ricardo Bustos M.
Jose Ricardo Bustos M.

Reputation: 8164

not much to explain, only the division and module is used to find the row and column number

items = [1,2,3,4,5,6,7,8,9]  

def SplitSets(iterable,set_size): 
  #return a generator
  return ((i/set_size, i%set_size,e) for i,e in enumerate(iterable))

for a,b,c in SplitSets(items,3):
  print a,b,c

a similar code is:

items = [1,2,3,4,5,6,7,8,9]

def SplitSets(iterable,set_size):
    for i, elem in enumerate(iterable):
        yield i/set_size, i%set_size, elem

for a,b,c in SplitSets(items,3):
    print a,b,c

Upvotes: 4

Related Questions