itertools in python

itertools.product()¶

Find out Cartesian product of input iterables.

In [64]:

#help('itertools.product')

In [9]:

from itertools import product

a = list(map(int, input().split()))
b = list(map(int, input().split()))

c = list(product(a, b))
print(' '.join([str(i) for i in c]))

1 2
3 4
(1, 3) (1, 4) (2, 3) (2, 4)

itertools.permutations()¶

Return successive r-length permutations of elements in the iterable.

In [66]:

#help('itertools.permutations')

In [49]:

from itertools import permutations

string, k = input().split()

p = list(permutations(string, int(k)))
print(p)
print('\n'.join([''.join(i) for i in sorted(p)]))

hack 2
[('h', 'a'), ('h', 'c'), ('h', 'k'), ('a', 'h'), ('a', 'c'), ('a', 'k'), ('c', 'h'), ('c', 'a'), ('c', 'k'), ('k', 'h'), ('k', 'a'), ('k', 'c')]
ac
ah
ak
ca
ch
ck
ha
hc
hk
ka
kc
kh

itertools.combinations()¶

Return successive r-length combinations of elements in the iterable.

In [68]:

#help('itertools.combinations')

In [34]:

from itertools import combinations

string, n = input().split()

p = [list(combinations(sorted(string), k)) for k in range(1, int(n)+1)]
print('\n'.join([''.join(j) for i in p for j in i]))

hack 2
a
c
h
k
ac
ah
ak
ch
ck
hk

itertools.combinations_with_replacement()¶

Return successive r-length combinations of elements in the iterable allowing individual elements to have successive repeats.

In [70]:

#help('itertools.combinations_with_replacement')

In [35]:

from itertools import combinations_with_replacement

string, n = input().split()

p = list(combinations_with_replacement(sorted(string), int(n)))
print('\n'.join([''.join(i) for i in p]))

hack 2
aa
ac
ah
ak
cc
ch
ck
hh
hk
kk

Iterables and Iterators¶

In [54]:

from itertools import combinations

n = int(input())
letters = input().split()
k = int(input())

perm = list(combinations(letters, k))

res = [l for l in perm if 'a' in l]

print(len(res)/len(perm))

4
a a c d
2
0.8333333333333334

Compress the String!¶

Make an iterator that returns consecutive keys and groups from the iterable.

In [72]:

#help('itertools.groupby')

In [62]:

from itertools import groupby

s = input()

result = ' '.join([str((len(list(k)), int(g))) for g, k in groupby(s)])

print(result)

1222311
(1, 1) (3, 2) (1, 3) (2, 1)

Maximize It!¶

In [48]:

from itertools import product

k, m = list(map(int, input().split()))

# as the first element indicates the length of input, [1:] is used to give up that
num = [list(map(int, input().split()))[1:] for i in range(k)]

result = map(lambda x: sum(i**2 for i in x)%m, product(*num))

print(max(result))
    

3 1000
2 5 4
3 7 8 9
5 5 7 8 9 10
206

In [ ]: