From 544f23bcd34a48f06ded2b3463a5ffd7dac44429 Mon Sep 17 00:00:00 2001 From: Raymond Hettinger Date: Thu, 11 Jan 2018 22:48:38 -0800 Subject: [PATCH 1/3] Add nth_combination recipe and test --- Doc/library/itertools.rst | 22 ++++++++++++++++++++++ Lib/test/test_itertools.py | 29 +++++++++++++++++++++++++++++ 2 files changed, 51 insertions(+) diff --git a/Doc/library/itertools.rst b/Doc/library/itertools.rst index 5af5422eefee67..59fdc86955f0f5 100644 --- a/Doc/library/itertools.rst +++ b/Doc/library/itertools.rst @@ -859,6 +859,28 @@ which incur interpreter overhead. indices = sorted(random.randrange(n) for i in range(r)) return tuple(pool[i] for i in indices) + def nth_combination(population, r, index): + 'Equivalent to list(combinations(population, r))[index]' + n = len(population) + c = 1 + k = min(r, n-r) + for i in range(1, k+1): + c = c * (n - k + i) // i + if index < 0: + index += c + if index < 0 or index >= c: + raise IndexError + if r < 0 or r > n: + raise ValueError + result = [] + while r: + c, n, r = c*r//n, n-1, r-1 + while index >= c: + index -= c + c, n = c*(n-r)//n, n-1 + result.append(population[-1-n]) + return tuple(result) + Note, many of the above recipes can be optimized by replacing global lookups with local variables defined as default values. For example, the *dotproduct* recipe can be written as:: diff --git a/Lib/test/test_itertools.py b/Lib/test/test_itertools.py index a359905e2e48cb..d3f53bc274dcb6 100644 --- a/Lib/test/test_itertools.py +++ b/Lib/test/test_itertools.py @@ -2262,6 +2262,29 @@ def test_permutations_sizeof(self): ... # first_true([a,b], x, f) --> a if f(a) else b if f(b) else x ... return next(filter(pred, iterable), default) +>>> def nth_combination(population, r, index): +... 'Equivalent to list(combinations(population, r))[index]' +... n = len(population) +... c = 1 +... k = min(r, n-r) +... for i in range(1, k+1): +... c = c * (n - k + i) // i +... if index < 0: +... index += c +... if index < 0 or index >= c: +... raise IndexError +... if r < 0 or r > n: +... raise ValueError +... result = [] +... while r: +... c, n, r = c*r//n, n-1, r-1 +... while index >= c: +... index -= c +... c, n = c*(n-r)//n, n-1 +... result.append(population[-1-n]) +... return tuple(result) + + This is not part of the examples but it tests to make sure the definitions perform as purported. @@ -2345,6 +2368,12 @@ def test_permutations_sizeof(self): >>> first_true('ABC0DEF1', '9', str.isdigit) '0' +>>> population = 'ABCDEFGH' +>>> for r in range(len(population) + 1): +... for i, expected in enumerate(combinations(population, r)): +... actual = nth_combination(population, r, i) +... assert expected == actual + """ __test__ = {'libreftest' : libreftest} From b629ad9207f7a311342baeb1e989327d70ecaead Mon Sep 17 00:00:00 2001 From: Raymond Hettinger Date: Thu, 11 Jan 2018 22:55:59 -0800 Subject: [PATCH 2/3] Use same parameter names as the other recipes --- Doc/library/itertools.rst | 9 +++++---- Lib/test/test_itertools.py | 9 +++++---- 2 files changed, 10 insertions(+), 8 deletions(-) diff --git a/Doc/library/itertools.rst b/Doc/library/itertools.rst index 59fdc86955f0f5..bb5d6dda38c2e3 100644 --- a/Doc/library/itertools.rst +++ b/Doc/library/itertools.rst @@ -859,9 +859,10 @@ which incur interpreter overhead. indices = sorted(random.randrange(n) for i in range(r)) return tuple(pool[i] for i in indices) - def nth_combination(population, r, index): - 'Equivalent to list(combinations(population, r))[index]' - n = len(population) + def nth_combination(iterable, r, index): + 'Equivalent to list(combinations(iterable, r))[index]' + pool = tuple(iterable) + n = len(pool) c = 1 k = min(r, n-r) for i in range(1, k+1): @@ -878,7 +879,7 @@ which incur interpreter overhead. while index >= c: index -= c c, n = c*(n-r)//n, n-1 - result.append(population[-1-n]) + result.append(pool[-1-n]) return tuple(result) Note, many of the above recipes can be optimized by replacing global lookups diff --git a/Lib/test/test_itertools.py b/Lib/test/test_itertools.py index d3f53bc274dcb6..4167f6ab1d5a99 100644 --- a/Lib/test/test_itertools.py +++ b/Lib/test/test_itertools.py @@ -2262,9 +2262,10 @@ def test_permutations_sizeof(self): ... # first_true([a,b], x, f) --> a if f(a) else b if f(b) else x ... return next(filter(pred, iterable), default) ->>> def nth_combination(population, r, index): -... 'Equivalent to list(combinations(population, r))[index]' -... n = len(population) +>>> def nth_combination(iterable, r, index): +... 'Equivalent to list(combinations(iterable, r))[index]' +... pool = tuple(iterable) +... n = len(pool) ... c = 1 ... k = min(r, n-r) ... for i in range(1, k+1): @@ -2281,7 +2282,7 @@ def test_permutations_sizeof(self): ... while index >= c: ... index -= c ... c, n = c*(n-r)//n, n-1 -... result.append(population[-1-n]) +... result.append(pool[-1-n]) ... return tuple(result) From da2746caf25d5f2571b8915cf86e1d5190e6ffd4 Mon Sep 17 00:00:00 2001 From: Raymond Hettinger Date: Fri, 12 Jan 2018 22:27:20 -0800 Subject: [PATCH 3/3] Move range check to start of function. Expand tests to include negative indicies. --- Doc/library/itertools.rst | 4 ++-- Lib/test/test_itertools.py | 13 ++++++++----- 2 files changed, 10 insertions(+), 7 deletions(-) diff --git a/Doc/library/itertools.rst b/Doc/library/itertools.rst index bb5d6dda38c2e3..0b3829f19faf9c 100644 --- a/Doc/library/itertools.rst +++ b/Doc/library/itertools.rst @@ -863,6 +863,8 @@ which incur interpreter overhead. 'Equivalent to list(combinations(iterable, r))[index]' pool = tuple(iterable) n = len(pool) + if r < 0 or r > n: + raise ValueError c = 1 k = min(r, n-r) for i in range(1, k+1): @@ -871,8 +873,6 @@ which incur interpreter overhead. index += c if index < 0 or index >= c: raise IndexError - if r < 0 or r > n: - raise ValueError result = [] while r: c, n, r = c*r//n, n-1, r-1 diff --git a/Lib/test/test_itertools.py b/Lib/test/test_itertools.py index 4167f6ab1d5a99..4fcc02acbf441c 100644 --- a/Lib/test/test_itertools.py +++ b/Lib/test/test_itertools.py @@ -2266,6 +2266,8 @@ def test_permutations_sizeof(self): ... 'Equivalent to list(combinations(iterable, r))[index]' ... pool = tuple(iterable) ... n = len(pool) +... if r < 0 or r > n: +... raise ValueError ... c = 1 ... k = min(r, n-r) ... for i in range(1, k+1): @@ -2274,8 +2276,6 @@ def test_permutations_sizeof(self): ... index += c ... if index < 0 or index >= c: ... raise IndexError -... if r < 0 or r > n: -... raise ValueError ... result = [] ... while r: ... c, n, r = c*r//n, n-1, r-1 @@ -2371,9 +2371,12 @@ def test_permutations_sizeof(self): >>> population = 'ABCDEFGH' >>> for r in range(len(population) + 1): -... for i, expected in enumerate(combinations(population, r)): -... actual = nth_combination(population, r, i) -... assert expected == actual +... seq = list(combinations(population, r)) +... for i in range(len(seq)): +... assert nth_combination(population, r, i) == seq[i] +... for i in range(-len(seq), 0): +... assert nth_combination(population, r, i) == seq[i] + """