Built-in tools for fast extraction of tuples of a given grade with respect to grading on [1, 2, ..., n]

edit

For decent values of $n$ one can try the following:

N = 6
S = [1..N]
grading = {1: 1, 2:1, 3: 1, 4: 2, 5: 2, 6: 3}

# Question 0
S2 = Combinations(S,2) 
S2_list = S2.list()
# Alternatively:
S2_list = [(j, k) for j in [1..N] for k in [j+1..N]]

# Question 1
g0 = 4
S2_list_in_deg_g0 = [(j, k) for j in [1..N] for k in [j+1..N] if grading[j] + grading[k] == g0]

# Question 2 :: same game
S3_list_in_deg_g0 = \
    [(j, k, n) for j in [1..N] for k in [j+1..N] for n in [k+1..N] 
     if grading[j] + grading[k] + grading[n] == g0]

S3_list_in_deg_g0_starting_with_one = \
    [(1, k, n) for k in [2..N] for n in [k+1..N] 
     if 1 + grading[k] + grading[n] == g0]

This is slow, since the python loops are slow, but should work good enough for values of $N$ up to $10000$ say. But i wonder if this lists are really needed as lists, to be stored on the HD as such. Instead, if the lists are an intermediate step, iterables may be better. (But python loops are still to slow.) To build the last as an iterable...

S3_iterable_in_deg_g0_starting_with_one = \
    ((1, k, n) for k in [2..N] for n in [k+1..N] 
     if 1 + grading[k] + grading[n] == g0)

For instance, consuming the iterable:

sage: S3_iterable_in_deg_g0_starting_with_one = \ 
....:     ((1, k, n) for k in [2..N] for n in [k+1..N]  
....:      if 1 + grading[k] + grading[n] == g0) 
....:                                                                                                                                                                           
sage: S3_iterable_in_deg_g0_starting_with_one                                                                                                                                   
<generator object <genexpr> at 0x7f5f920eba50>
sage: for triple in S3_iterable_in_deg_g0_starting_with_one: 
....:     print(triple) 
....:                                                                                                                                                                           
(1, 2, 4)
(1, 2, 5)
(1, 3, 4)
(1, 3, 5)

The above solutions are using list comprehension. If the above is slow for the needed purpose, than to improve consider the following ideas:

• I suppose that $N$ may be bigger $10^6$, but the "grading(s)" have a small range. In our case, the values of the gradings / weights dictionary, taken as a set, is {1,2,3}. Suppose this list is relatively small. (Less than $10^5$.) Then build all tuples of two weights which sum up to g0. Then define the iterable by passing through the weights first.

For instance:

weights = list(set(grading.values()))
weights2 = [(u, v) for u in weights for v in weights if u <= v if u + v == g0]

S2_iterable_in_deg_g0 = \
        ((j, k) for (u, v) in weights2 for j in [1..N-1] for k in [j+1..N] 
         if grading[j] == u and grading[k] == v)

# consume it via...
for tup in S2_iterable_in_deg_g0:
    print(tup)

(The above can be done better...)

• Use numpy for vectorialization.

There are two functions in Sage which are similar, but do not do exactly what you need. The implementations might provide useful ideas though. This one does not respect the grading, but otherwise solves the second problem. It is implemented in Cython, so is probably quite fast:

sage: IntegerListsLex(n=9, length=3, min_slope=1, min_part=1, max_part=6).list()
[[2, 3, 4], [1, 3, 5], [1, 2, 6]]

This one respects the grading. Though, it does not enumerate subsets, but multisets (subsets with repetition):

sage: WeightedIntegerVectors(4, (1,1,1,2,2,3)).list()
[[0, 0, 1, 0, 0, 1],
 [0, 1, 0, 0, 0, 1],
 [1, 0, 0, 0, 0, 1],
 [0, 0, 0, 0, 2, 0],
 ...

In general, you will want to avoid constructing elements of the wrong grade, as early as possible.