List comprehension in Mathematica

Python supports a concept called "list comprehensions".

Following is sample example (written in Python syntax)

A = [[1, 2, 3], [2, 3, 4], [3, 4, 5]]
print( [ a + [i]  for a in A for i in range(1, 4) if i in a] )

[[1, 2, 3, 1], [1, 2, 3, 2], [1, 2, 3, 3], [2, 3, 4, 2], [2, 3, 4, 3], [3, 4, 5, 3]]


What's the most natural way to translate this Python code to Mathematica? I have the following method:

A = {{1, 2, 3}, {2, 3, 4}, {3, 4, 5}};
Select[Table[ a ~ Join ~ {i}, {a, A}, {i, 3}] ~ Flatten ~ 1,
MemberQ[ Most @ #, Last @ #]&]


but I think it's neither fast nor elegant, therefore I am looking for a better way.

• The same article you linked has a few lines about Mathematica too. Commented Nov 28, 2013 at 12:05
• Maybe this? Commented Nov 28, 2013 at 13:03

This one is very close to your Python code

Join @@ Table[Append[a, i], {a, A}, {i, Intersection[Range[3], a]}]

{{1, 2, 3, 1}, {1, 2, 3, 2}, {1, 2, 3, 3}, {2, 3, 4, 2}, {2, 3, 4, 3}, {3, 4, 5, 3}}

• After I overviewed this answer here, I acknowledged Mathematica's potential to build new styles -- instead of using array structures we could define {a,A} to be "for a in A" and "... if a in A" could be just the intersection. I think we could extend your answer here to create a visually-accessible Pythonic style to Mathematica. What do you think?
– hhh
Commented Jun 27, 2015 at 22:28
• This assumes that one can compute the set of valid a ahead of time, and then to implement some way to do that computation before performing the operation. A list comprehension is far more general. Commented Apr 28, 2020 at 0:53
A = {{1, 2, 3}, {2, 3, 4}, {3, 4, 5}};
Join @@ Table[If[a~MemberQ~i, a~Join~{i}, Unevaluated[]], {a, A}, {i, 3}]
Join @@ Table[a~Join~{i}, {a, A}, {i, Select[Range@3, a~MemberQ~# &]}]


{{1, 2, 3, 1}, {1, 2, 3, 2}, {1, 2, 3, 3}, {2, 3, 4, 2}, {2, 3, 4, 3}, {3, 4, 5, 3}}

• +1 for writing Unevaluated[] instead of Unevaluated[Sequence[]], or ##&[] :) Commented Nov 28, 2013 at 14:32
• @Jacob What's wrong with my "vanishing function" ##&[]? (+1, by the way) Commented Nov 28, 2013 at 16:30
• @Mr.Wizard haha nothing really, in fact I was glad to show off that I knew about it :P. But also despite all the talk about this, I may have missed the option of just using Unevaluated[], so I guess I learned something :). Commented Nov 28, 2013 at 16:49
• @Jacob I attempted to cover that in this answer but I guess I should have been more explicit. I said that "All of these heads meet this requirement:" and listed Unevaluated and ## & -- both would be used the same way: Unevaluated[] and ## &[]. Commented Nov 29, 2013 at 12:20
• @Mr.Wizard that one is a nice answer (as I suppose you can tell by the badge :P). I'm not very focussed right now, but perhaps an additional point of confusion could by that neither Apply[] nor Symbol[] evaluate to Sequence[] (anymore?). Commented Nov 29, 2013 at 12:42

Some of the other approaches might be much more efficient, but the following shows how one can create something which is probably as easy to read (if one is fluent in Mathematica) as a python list comprehension:

SetAttributes[listComprehend, HoldAll]

listComprehend[Verbatim[Condition][body_, crit_],iters:({_, __}..)] := Flatten[
Table[
If[crit, body, Unevaluated[Sequence[]]],
iterators
], 1]


use it like this:

listComprehend[Append[a, i] /; MemberQ[a, i], {a, A}, {i, 3}]


/; is the shortcut for Condition and usually can be read as "provided that" in mathematica code. Doing so, the above seems to be relatively clear code. Unfortunately the translation to one of the potentially more efficient approaches is not that simple...

Not really sure if you're after the functionality, the way to generate your list, or some performance improvement. Anyway:

Cases[Outer[Join, A, List /@ Range@Length @A, 1], {___, x_, ___, x_}, 2]
(*
{{1, 2, 3, 1}, {1, 2, 3, 2}, {1, 2, 3, 3},
{2, 3, 4, 2}, {2, 3, 4, 3}, {3, 4, 5, 3}}
*)


Edit

For efficiency, you may want to generate only the final elements, without the need to post filter them:

Flatten[Function[{u}, Table[Join[u, {i}], {i, u}]][#] & /@ A, 1]


or:

Flatten[Function[{u}, ReplaceList[#, Thread[x_ :> Flatten[{x, #}] & /@ u]]][#] & /@ A, 1]


If you are going for a literal translation from Python to Mathematica, I think that something like the following snippet should do the trick:

Part[
Reap@Do[
Do[
If[
MemberQ[a, r],
Sow@Flatten[{a, r}]
], {r, R}], {a, A}], 2, 1]


where A = {{1, 2, 3}, {2, 3, 4}, {3, 4, 5}} and R=Range[1,3] (note how in Python $list(range(1, 4))=[1, 2, 3]$)

Nevertheless, this question is about finding "...the most natural way to translate this Python code to Mathematica" and my view is that going for the literal translation is not always the best option (although it could very possibly be what you need at a certain point in time).

The answers to this question provide useful examples of how versatile Mathematica is.

Another natural way to achieve the intended result, is to use Map:

Flatten[
Map[(
a = #;
Flatten[{a, #}] & /@ Pick[R, MemberQ[a, #] & /@ R]
) &, A], 1]


Perhaps, a more concise way to achieve the same result is the following:

f[a_, r_] := If[MemberQ[a, r], Flatten[{a, r}]]

Cases[
Flatten[
Outer[f, A, R, 1], 1], Except[Null], 1]


I think that this way is more close in practice to what the list comprehension does: it traverses the items of A (rows) and then it traverses the items in the range until it finds a match.Then it appends the range item to the list row and proceeds to the next row.