Added generator for nth product and nth_products.

Author: root-11

graph/generate.py

@@ -1,3 +1,7 @@
+from functools import lru_cache
+from collections import defaultdict
+import math
+import random
 import itertools
 from graph.core import Graph
 
@@ -36,3 +40,84 @@ def grid(length, width, bidirectional=False):
                 a, b = node_index[(i, j)], node_index[(i, j - 1)]
                 g.add_edge(b, a, bidirectional=bidirectional)
     return g
+
+
+def nth_permutation(idx, length, alphabet=None, prefix=()):
+    if alphabet is None:
+        alphabet = [i for i in range(length)]
+    if length == 0:
+        return prefix
+    else:
+        branch_count = math.factorial(length - 1)
+        for d in alphabet:
+            if d not in prefix:
+                if branch_count <= idx:
+                    idx -= branch_count
+                else:
+                    return nth_permutation(idx, length - 1, alphabet, prefix + (d,))
+
+
+def nth_product(idx, *args):
+    """returns the nth product of the given iterables.
+
+    Args:
+        idx (int): the index.
+        *args: the iterables.
+    """
+    if not isinstance(idx, int):
+        raise TypeError(f"Expected int, not {type(idx)}")
+    total = math.prod([len(a) for a in args])
+    if idx < 0:
+        idx += total
+    if index < 0 or index >= total:
+        raise IndexError(f"Index {index} out of range")
+    
+    elements = ()
+    for i in range(len(args)):
+        offset = math.prod([len(a) for a in args[i:]]) // len(args[i])
+        index = idx // offset
+        elements += (args[i][index],)
+        idx -= index * offset
+    return elements
+
+
+def n_products(*args, n=20):
+    """
+    Returns the nth product of the given iterables.
+    """
+    if len(args) == 0:
+        return ()
+    if any(len(a) == 0 for a in args):
+        raise ZeroDivisionError("Cannot generate products of empty iterables")
+
+    n = min(n, math.prod([len(a) for a in args]))
+    step = math.prod([len(a) for a in args]) / n
+
+    for ni in range(n):
+        ix = int(step * ni + step / 2)
+        yield nth_product(ix, *args)
+    
+
+def random_graph(size, degree=1.7, seed=1):
+    if not isinstance(size, int):
+        raise TypeError(f"Expected int, not {type(size)}")
+    if not isinstance(degree, float):
+        raise TypeError(f"Expected float, not {type(degree)}")
+    if not isinstance(seed, int):
+        raise TypeError(f"Expected int, not {type(seed)}")
+
+    g = Graph()
+    nodes = list(range(size))
+    rng = random.Random(seed)
+    rng.shuffle(nodes)
+
+    edges = size * degree
+
+    gen_length = math.factorial(nodes) / math.factorial(nodes - 2)
+    comb = int(gen_length / edges)
+
+    for i in nodes:
+        n = i * comb
+        a, b = nth_permutation(n, size, nodes)
+        g.add_edge(a, b)
+    return g

tests/test_generate.py

@@ -1,5 +1,6 @@
+import math, itertools
 from graph.core import Graph
-from graph.generate import binary_tree, grid
+from graph.generate import binary_tree, grid, random_graph, n_products, nth_product
 
 from tests.test_graph import graph3x3, graph4x4, graph5x5
 
@@ -15,12 +16,43 @@ def test_generate_grid_4x4_bidirectional():
     expected = graph4x4()
     assert g == expected
 
+
 def test_generate_grid_5x5_bidirectional():
     g = grid(5, 5, bidirectional=True)
     expected = graph5x5()
     assert g == expected
 
+
 def test_generate_binary_tree_3_levels():
     g = binary_tree(3)
     expected = Graph(from_list=[[0, 1], [0, 2], [1, 3], [1, 4], [2, 5], [2, 6]])
-    assert g == expected
+    assert g == expected
+
+
+def test_nth_p():
+    a, b, c, d = [1, 2, 3, 4, 5, 6], [1, 2, 3, 4, 5], [1, 2, 3, 4], [1, 2, 3, 4, 5, 6, 7]
+    for i, comb in enumerate(itertools.product(*[a, b, c, d])):
+        if i < 5 or i > 836 or i == 444:
+            print(i, comb)
+        assert comb == nth_product(i, a, b, c, d)
+
+
+def test_nth_products():
+    a, b, c = [1, 2, 3, 4, 5, 6], [1, 2, 3, 4, 5, 6], [1, 2, 3, 4, 5, 6]
+    L = list(n_products(a, b, n=10))
+    assert len(L) == 10
+    assert L == [(1, 2), (1, 6), (2, 4), (3, 1), (3, 5), (4, 2), (4, 6), (5, 4), (6, 1), (6, 5)]
+
+    L = list(n_products(a, b, c, n=10))
+    assert len(L) == 10
+    # fmt:off
+    assert L == [(1, 2, 5), (1, 6, 3), (2, 4, 1), (3, 1, 4), (3, 5, 2), (4, 2, 5), (4, 6, 3), (5, 4, 1), (6, 1, 4), (6, 5, 2)]
+    # fmt:on
+
+
+def test_random_graphs():
+    g1 = random_graph(10, 1.7, 1)
+    g2 = random_graph(10, 1.7, 1)
+    assert g1 == g2
+    g3 = random_graph(10, 1.1, 1)
+    assert g1 != g3