List comprehension in Mathematica

Question

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.

Answer 1

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}}

Answer 2

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}}

Answer 3

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...

Answer 4

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]

Answer 5

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.