What general tips do you have for golfing in Python? I'm looking for ideas which can be applied to code-golf problems and which are also at least somewhat specific to Python (e.g. "remove comments" is not an answer).
Please post one tip per answer.
Code Golf Stack Exchange is a question and answer site for programming puzzle enthusiasts and code golfers. It only takes a minute to sign up.
Sign up to join this communityWhat general tips do you have for golfing in Python? I'm looking for ideas which can be applied to code-golf problems and which are also at least somewhat specific to Python (e.g. "remove comments" is not an answer).
Please post one tip per answer.
These are some golfed implementations of number theoretic functions that come up in challenges a lot. Many of these are due to xnor, especially the “Wilson’s theorem prime machines” of the form lambda n,i=1,p=1
. The coprime/totient functions are Dennis’s (explanation here).
It is instructive to study what exactly these are doing, so that you can adapt them to your needs or roll them into another recursive function. That often ends up being shorter than pasting these directly into your solution as-is!
All of these assume n
is a positive integer. The ones marked with an asterisk produce the wrong result if n = 1
. Furthermore, these snippets assume Python 2. For Python 3, you might need to replace /
by //
here and there.
# Function Output of f(360)
#========================================================================================
f=lambda n,i=2:n/i*[0]and[f(n,i+1),[i]+f(n/i)][n%i<1] # [2, 2, 2, 3, 3, 5] (slow!)
f=lambda n,i=2:n/i*[0]and f(n,i+1)if n%i else[i]+f(n/i) # [2, 2, 2, 3, 3, 5]
f=lambda n,i=2:n/i*[0]and(n%i and f(n,i+1)or[i]+f(n/i)) # [2, 2, 2, 3, 3, 5]
f=lambda n,i=2:n<2and{1}or n%i and f(n,i+1)or{i}|f(n/i) # {1, 2, 3, 5}
f=lambda n,i=2:n<2and{i}or n%i and f(n,i+1)or{i}|f(n/i,i) #*{2, 3, 5}
f=lambda n,i=2:n/i and[f(n,i+1),i+f(n/i)][n%i<1] # 2+2+2+3+3+5 (slow!)
f=lambda n,i=2:n/i and f(n,i+1)if n%i else i+f(n/i) # 2+2+2+3+3+5
f=lambda n,i=2:n/i and(n%i and f(n,i+1)or i+f(n/i)) # 2+2+2+3+3+5
f=lambda n,i=1,p=1:n*[0]and p%i*[i]+f(n-p%i,i+1,p*i*i) # first n primes
f=lambda n,i=1,p=1:n*[0]and p%i*[i]+f(n-1,i+1,p*i*i) # primes <= n
f=lambda n,i=1,p=1:n/i and p%i*i+f(n,i+1,p*i*i) # sum of primes <= n
f=lambda n,i=1,p=1:n/i and p%i+f(n,i+1,p*i*i) # count primes <= n
f=lambda n,i=1,p=1:n and-~f(n-p%i,i+1,p*i*i) # nth prime
f=lambda n:all(n%m for m in range(2,n)) #*is n prime? (not recursive)
f=lambda n:1>>n or n*f(n-1) # factorial
f=lambda n:sum(k/n*k%n>n-2for k in range(n*n)) # totient phi(n) (not recursive)
f=lambda n:[k/n for k in range(n*n)if k/n*k%n==1] # coprimes up to n (not recursive)
Additions and byte saves are very welcome!
)
. (Also it has a variant similar to #2.)
\$\endgroup\$
Jan 4, 2018 at 0:16
f=lambda n:1>>n or n*f(n-1)
be lambda n:n<2or n*f(n-1)
, or am I going crazy?
\$\endgroup\$
f(0) == f(1) == True
rather than 1
.
\$\endgroup\$
from math import perm as f
is shorter than factorial
definition. and that gives a shorter way to check primality using Wilson's theorem if math
is already imported (first used in codegolf.stackexchange.com/a/218738/69850, ls mentioned in codegolf.stackexchange.com/a/194319/69850)
\$\endgroup\$
Feb 8, 2021 at 0:38
Suppose you have a string s
, and need to print it without a trailing newline. The canonical way of doing this would be
print(s,end='')
However, if we look at the documentation for print
, we can see that print
takes in a variable number of objects as its first parameter, with "variable number" including zero. This means that we can do
print(end=s)
instead, for a saving of 3 bytes.
Note that this only works when s
is a string, since otherwise the conversion to string would be too expensive:
print(1,end='')
print(end=str(1))
Thanks to @Reticality for this tip.
``
commenters: ``
does not work in Python 3 (This is a spam prevention comment)
\$\endgroup\$
Jun 25, 2016 at 13:56
>>> for i in range(x):s+=input()
if value of i is useless:
>>> for i in[0]*x:s+=input()
or
>>> exec's+=input();'*x
for i in[0]*x:s+=input()
to save another space. Also, you can remove the space between the exec and the first quotation mark to get exec's+=input();'*x
\$\endgroup\$
Apr 19, 2011 at 6:12
for i in[0]*x:s+=input()
\$\endgroup\$
Aug 5, 2015 at 7:47
Multiple statements can be put on one line separated by ;
. This can save a lot of whitespace from indentation.
while foo(a):
print a;a*=2
Or even better:
while foo(a):print a;a*=2
if
s and while
s inside the while
.
\$\endgroup\$
Oct 29, 2011 at 20:23
Assignment expressions are a powerful language feature introduced in Python 3.8 (TIO). Use the "walrus operator" :=
to assign a variable inline as part of expression.
>>> (n:=2, n+1)
(2, 3)
You can save an expression to a variable inside a lambda
, where assignments are not ordinarily allowed. Compare:
def f(s):t=s.strip();return t+t[::-1]
lambda s:s.strip()+s.strip()[::-1]
lambda s:(t:=s.strip())+t[::-1]
An assignment expression can be used in a comprehension to iteratively update a value, storing the result after each step in a list or other collection. This example computes a running sum by updating the running total t
.
>>> t=0
>>> l=[1,2,3]
>>> print([t:=t+x for x in l])
[1, 3, 6]
>>> t
6
This can be done in a lambda
with the initial value as an optional argument:
>>> f=lambda l,t=0:[t:=t+x for x in l]
>>> f([1,2,3])
[1, 3, 6]
This function is reusable: each call with start t
back at 0.
Here are some arithmetic tricks which are either shorter or are more useful due to precedence rules.
Assumptions Version 1 Version 2
-------------------------------------------------------------------
n >= 0 float n==0 0**n
n >= 0 integer n==0 1>>n
n > 0 integer n!=1 1%n
n > 0 integer, Python 2 n==1 1/n
n, m float n!=m n-m
1//n
could still be useful in Python 3 for n==1
since they have different precedence.
\$\endgroup\$
You may know the trick to easily create an infinite generator using the two-argument form of iter
:
(... for _ in iter(lambda:0,1))
where 0
and 1
can be any two non-equal values.
As @ovs once pointed out to me, you can replace lambda:0
with int
, because int
returns 0 when called with no arguments.
(... for _ in iter(int,1))
However, we can do one byte better!:
(... for()in iter(set,1))
This uses the fact that ()
is a valid L-value (assignment target) in Python as long as the RHS is an empty sequence, and set
returns an empty set when called with no arguments.
You can also use str
instead of set
for the same byte count.
I've seen this situation pop up a few times, so I thought a tip would be good.
Suppose you have a string s
and you want to translate some chars of s
to other chars (think ROT-13 like ciphers). For a more concrete example, suppose we want to swap just the a
s and b
s in a string, e.g.
"abacus" -> "babcus"
The naïve way to do this would be:
lambda s:s.replace('a','T').replace('b','a').replace('T','b')
Note how we need to introduce a temporary 'T'
to get the swapping right.
With eval
, we can shorten this a bit:
lambda s:eval("s"+".replace('%s','%s')"*3%tuple("aTbaTb"))
For this particular example, iterating char-by-char gives a slightly better solution (feel free to try it!). But even so, the winner is str.translate
, which takes a dictionary of from: to
code points:
# Note: 97 is 'a' and 98 is 'b'
lambda s:s.translate({97:98,98:97})
In Python 2 this only works for Unicode strings, so unfortunately the code here is slightly longer:
lambda s:(u''+s).translate({97:98,98:97})
Some important points which make str.translate
so useful are:
"cus"
in "abacus"
above.to
part of the dictionary can actually be a (Unicode) string as well, e.g. {97:"XYZ"}
(u"XYZ"
in Python 2) would turn abacus -> XYZbXYZcus
. It can also be None
, but that doesn't save any bytes compared to ""
or u""
."ab".translate({97:None})
is longer than "ab".translate({97:""})
.
\$\endgroup\$
Oct 8, 2015 at 8:16
from string import*;t="abcdefghijklmABCDEFGHIJKLMZYXWVUTSRQPONzyxwvutsrqpon";lambda s:s.translate(maketrans(t,t[::-1]))
\$\endgroup\$
Aug 18, 2017 at 4:05
The binomial coefficient \$\binom{n}{k} \ = \frac{n!}{k!(n-k)!}\$ can be expressed arithmetically as
((2**n+1)**n>>n*k)%2**n
This works for \$n,k \geq 0\$, except for \$n=k=0\$ it gives \$0\$ rather than \$1\$. More generally, it works to use
(b+1)**n/b**k%b
(TIO), where \$b\$ is any value strictly greater than the result. The first expression uses \$b=2^n\$, which exceeds \$\binom{n}{k}\$ except for \$n=k=0\$.
Why does this work? Let's look at an example with b=1000
. Then, for n=6
, we have
(b+1)**n = 1001 ** 6 = 1006015020015006001
Note how triples of digits encode the binomial coefficients in the n=6
row of Pascal's triangle:
1 6 15 20 15 6 1
1 006 015 020 015 006 001
This works because the binomial coefficients are the coefficients of the polynomial
$$ (b+1)^n = \sum_{k=0}^n\binom{n}{k}b^k$$
and so can be read off as digits in base b
, as long no binomial coefficient exceeds b
which would cause regrouping.
We can extract a given triple of digits, say for \$\binom{6}{2}=15\$, by floor-dividing by 1000000
to delete the last 6 digits leaving 1006015020015
, then take %1000
to extract the last triplet 015
. More generally, doing /b**k%b
extracts the k
-th digit from the end zero-indexed in base b
, that is the digit with multiplier b**k
.
If you're doing somewhat more complex golfing that require something from the standard library to be used a lot, import x as y
can save some space:
import itertools as i
i.groupby(...) # same as itertools.groupby
itertools
are just way too descriptive, makes me sad.
\$\endgroup\$
In Python 3, the built-in function open
underwent some changes. In particular, its first argument
file is either a string or bytes object giving the pathname (absolute or relative to the current working directory) of the file to be opened or an integer file descriptor of the file to be wrapped.
(source)
That means
open(0).read()
suffices to read all input from STDIN.
Try it online on Ideone.
OSError: [WinError 6] The handle is invalid
)
\$\endgroup\$
So you have a Boolean... a real Boolean, not one represented as an integer. You have a condition where it needs to be negated, and you can't just go back and negate it where you got it (e.g. !=
instead of ==
), maybe because you use it once straight and once negated.
Well, who says your Booleans aren't longing to be integers deep in their little hearts?
>>> False < 1
True
>>> True < 1
False
8 bytes, not counting the colon:
if not C:
6 bytes:
if C<1:
EDIT: 5 bytes, thanks to user202729 in the comments:
if~-C:
This works because:
>>> -False
0
>>> -True
-1
>>> ~-False
-1
>>> ~-True
0
You can split a list into chunks of a given size using zip
and iter
, as explained in this SO question.
>>> l=range(12)
>>> zip(*[iter(l)]*4)
[(0, 1, 2, 3), (4, 5, 6, 7), (8, 9, 10, 11)]
Of course, substituting in l
as zip(*[iter(range(12))]*4)
gives the same result.
The 4
is the number of elements per chunk. If the length isn't a multiple of this, any elements in the remainder are not included. For example, l=range(13)
would give the same result.
The result is a list of tuples. If your input is a string and you want to produce a list of strings, you can do
>>> l="Code_golf"
>>> map(''.join,zip(*[iter(l)]*3))
['Cod', 'e_g', 'olf'] # Python 3 would give a map object
When the list l
is defined by a list comprehension, instead of converting to an iterable as iter(l)
, you can instead write it as a generator comprehension with (...)
instead of [...]
.
>>> l=(n for n in range(18)if n%3!=1)
>>> zip(*[l]*4)
[(0, 2, 3, 5), (6, 8, 9, 11), (12, 14, 15, 17)]
This consumes the generator, so l
will appear empty afterwards. Note as before that we can inline l
as zip(*[(n for n in range(18)if n%3!=1)]*4)
.
To replace every entry of value a
with b
in a list L
, use:
map({a:b}.get,L,L)
For example,
L=[1,2,3,1,2,3]
a=2
b=3
print map({a:b}.get,L,L)
[1, 3, 3, 1, 3, 3] #Output
In Python 3, this returns a map
object rather than a list. The list entries can be any hashable values (ints, floats, strings, tuples, etc).
Here's how this works. A dictionary's get
method takes a key and default value, and returns the dictionary's entry for that key, using the default value is the key is not present. This method is mapped method over each entry in L
both as the key and the default value, which results in
[{a:b}.get(x,x) for x in L]
If x
is a
, then the dictionary transforms it to b
, and otherwise, it defaults to itself. You can perform multiple replacements at the same time using a larger dictionary.
Credit to twobit on Anarchy Golf for exposing me to this trick.
You can check if a possibly-empty list l
starts with a value x
by doing
l[:1]==[x]
This gives False
on an empty list, while l[0]==x
gives an out-of-bounds error. Strings works similarly
s[:1]=='a'
In general, you can safely check the n
'th element as
l[n:n+1]==[a]
or as l[n:][:1]==[a]
when n
is a long expression.
map
can take multiple iterable arguments and apply the function in parallel.
Instead of
a=[1,4,2,6,4]
b=[2,3,1,8,2]
map(lambda x,y:...,zip(a,b))
you can write
map(lambda x,y:...,a,b)
You only need to indent nested control structures:
def baz(i):
if i==0:i=1;print i;bar()
while i:i+=foo(i-1)
You can use default arguments of a function to save some indentation, since
def f(a,l=[1,2,3]):
return sum(a==i for i in l)
is one byte shorter than
def f(a):
l=[1,2,3]
return sum(a==i for i in l)
pop
or l[0] =1
), that list will be changed in the outer scope too.
\$\endgroup\$
None
arguments in Python builtinsMapping with None
in place of a function assumes the identity function instead. This allows it to be used as an alternative to itertools.izip_longest
for zipping lists to the length of the longest list:
>>> L = [[1, 2], [3, 4, 5, 6], [7]]
>>> map(None,*L)
[(1, 3, 7), (2, 4, None), (None, 5, None), (None, 6, None)]
For visualisation (with .
representing None
):
1 2 1 3 7
3 4 5 6 -> 2 4 .
7 . 5 .
. 6 .
filter
with None
also assumes the identity function, thus removing falsy elements.
>>> L = ["", 1, 0, [5], [], None, (), (4, 2)]
>>> filter(None, L)
[1, [5], (4, 2)]
This is a bit better than a list comprehension:
filter(None,L)
[x for x in L if x]
However, as @KSab notes, if all elements are of the same type then there may be shorter alternatives, e.g. filter(str,L)
if all elements are strings.
filter(str,L)
if L is all strings or filter(int,L)
if all ints which in some cases could be shorter.
\$\endgroup\$
Use os.read
to read all input:
import os
s=os.read(0,1e9)
Which is shorter than
import sys
s=sys.stdin.read()
Note that this has a limitation on input length, but it's so ridiculously large I'd say we're safe from the angry mob.
raw_input()
is shorter. If you need to read once, just spelling it out is shorter than the import + os.read; if more than once, assign it to a single-character value.
\$\endgroup\$
Use extended slicing to select one of two strings
>>> for x in-2,2:print"WoolrlledH"[::x]
...
Hello
World
vs
>>> for x in 0,1:print["Hello","World"][x]
...
Hello
World
Just found out two new things. First, input()
can parse tuples, like 1, 2, 3
is equivalent to the tuple (1, 2, 3)
.
And if you need to convert a value to float, just multiply by 1.
. Yes, 1.
is valid syntax (At least in 2.6).
input(x)
is basically the same thing as eval(raw_input(x))
. Unsafe to use in practice, but good for code golfing.
\$\endgroup\$
Jan 24, 2014 at 2:33
input()
only once. E.g. a,b,c=input()
will read in three comma-separated arguments and assign them to a, b, and c
\$\endgroup\$
Feb 1, 2016 at 4:17
1
to save 1 byte?, because division always gives a float
\$\endgroup\$
Adding vectors
Python doesn't have a built-in way to do vector (component-wise) addition except with libraries. Say a
and b
are two equal-length lists of numbers you want to add. Instead of the list comprehension
c=[a[i]+b[i]for i in range(len(a))]
you can use
c=map(sum,zip(a,b))
This produces an annoying map
object in Python 3, but it's shorter even if you have to convert to a list.
map(int.__add__,a,b)
is more readable with 1 char longer.
\$\endgroup\$
If you want to know the type of a variable x:
x*0is 0 # -> integer
x*0==0 # -> float (if the previous check fails)
x*0=="" # -> string
x*0==[] # -> array
If you represent boolean values as numbers you can save characters. This is especially true for using -1
as True
.
Bitty conditionals work (Truth table):
a b & | ^
0 0 0 0 0
0 -1 0 -1 -1
-1 0 0 -1 -1
-1 -1 -1 -1 0
And ~
works as not:
a ~a
0 -1
-1 0
Even though the -
for initializing -1
costs one character, this can easily save characters overall.
Compare:
while~a:
to:
while not a:
.center
in ASCII artIn drawing a symmetrical ASCII art, you can center-justify each line in a fixed width of spaces. For example, "<|>".center(7)
gives ' <|> '
. This can be shorter than computing how many spaces are needed to center it.
You can also pad with a different character by doing "<|>".center(7,'_')
f"{'<|>':^7}"
not be shorter in python3.6+ for non-variable widths? Even more so when providing the character to center by, f"{'<|>':-^7}"
vs "<|>".center(7,"-")
\$\endgroup\$
Sep 13, 2019 at 10:48
Iterating over indices in a list
Sometimes, you need to iterate over the indices of a list l
in order to do something for each element that depends on its index. The obvious way is a clunky expression:
# 38 chars
for i in range(len(l)):DoStuff(i,l[i])
The Pythonic solution is to use enumerate
:
# 36 chars
for i,x in enumerate(l):DoStuff(i,x)
But that nine-letter method is just too long for golfing.
Instead, just manually track the index yourself while iterating over the list.
# 32 chars
i=0
for x in l:DoStuff(i,x);i+=1
Here's some alternatives that are longer but might be situationally better
# 36 chars
# Consumes list
i=0
while l:DoStuff(i,l.pop(0));i+=1
# 36 chars
i=0
while l[i:]:DoStuff(i,l[i]);i+=1
Combining with this tip, suppose you have a situation like
for _ in[0]*x:doSomething()
a="blah"
You can instead do:
for a in["blah"]*x:doSomething()
to skip out on a variable assignment. However, be aware that
exec"doSomething();"*x;a="blah"
in Python 2 is just shorter, so this only really saves in cases like assigning a char (via "c"*x
) or in Python 3.
However, where things get fun is with Python 2 list comprehensions, where this idea still works due to a quirk with list comprehension scope:
[doSomething()for a in["blah"]*x]
(Credits to @xnor for expanding the former, and @Lembik for teaching me about the latter)
If you want to make a 3*4 grid of zeroes, the natural expression M=[[0]*4]*3
gives an unpleasant surprise if you modify an entry:
>>> M=[[0]*4]*3
>>> M
[[0, 0, 0, 0], [0, 0, 0, 0], [0, 0, 0, 0]]
>>> M[0][0]=1
>>> M
[[1, 0, 0, 0], [1, 0, 0, 0], [1, 0, 0, 0]]
Since each row is a copy of the same list by reference, modifying one row modifies all of them, which is usually not the behavior you want.
In Python 2, avoid this with the hack (19 chars):
M=eval(`[[0]*4]*3`)
Doing eval(`_`)
converts to the string representation, then re-evaluates it, converting the object to the code of how it's displayed. In effect, it's doing copy.deepcopy
.
If you're OK getting a tuple of lists, you can do (18 chars):
M=eval('[0]*4,'*3)
to get ([0, 0, 0, 0], [0, 0, 0, 0], [0, 0, 0, 0])
. This lets you do M[0][0]=1
but not M[0]=[1,2,3,4]
. It also works in Python 3.
str
for backticks as M=eval(str([[0]*3]*4))
. Or, M=[3*[0]for _ in[0]*4]
which is the same length. Maybe there's better.
\$\endgroup\$
Say you have two 2-element tuples which represent points in the Euclidean plane, e.g. x=(0, 0)
and y=(3, 4)
, and you want to find the distance between them. The naïve way to do this is
d=((x[0]-y[0])**2+(x[1]-y[1])**2)**.5
Using complex numbers, this becomes:
c=complex;d=abs(c(*x)-c(*y))
If you have access to each coordinate individually, say a=0, b=0, c=3, d=4
, then
abs(a-c+(b-d)*1j)
can be used instead.
d=abs(x[0]-y[0]+(x[1]-y[1])*1j)
. Can this be made shorter than what you have? I'm just wondering if 1j
could be useful in converting a tuple to a complex number.
\$\endgroup\$
*1j
is usually shorter.
\$\endgroup\$
x=(0,0);y=(3,4);c=complex;d=abs(c(*x)-c(*y))
, different assignments: a,b=0,0;c,d=3,4;d=((a-c)**2+(b-d)**2)**.5
. Yes, writing x=0,0;y=3,4
does indeed make the complex option shorter, but using different assignments makes it even shorter: x=0,0;y=3,4;c=complex;d=abs(c(*x)-c(*y)
. Finally, consider @mbomb007's approach with different assignments: a,b=0,0;x,y=3,4;d=abs(a-x+(b-y)*1j)
: with or without the assignment, it's shorter than all of the alternatives I've found.
\$\endgroup\$
1j
part, so I've added that in. The general point was basically that ((a-b)**2+(c-d)**2)**.5
is rarely ever needed.
\$\endgroup\$
:=
operator in 3.8 \$\endgroup\$