Tips for golfing in Python

Question

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.

Answer 1

To convert a 2D list to a 1D flat list, you can do sum(2d_list,[]) instead of the long way of import numpy;list=list.flatten(), which requires to create a numpy.array at the start as well, which is way longer than the first one

Answer 2

If you have multidimensional array of numbers and for instance need to count all numbers greater than n.

First flatten the array, then apply filter function to match condition:

l=[[1,[8,4,7,1],3],[5,[7],3,9],[7,3,9,[[[8]]]]]
n=5
flatten=lambda l: sum(map(flatten,l),[]) if isinstance(l,list) else [l]
len(filter(lambda x:x>n,flatten(l)))

Answer 3

To find the all the indexes of a certain element in a list l, use

filter(lambda x:l[x]==element,range(len(l)))

To find the next index after a certain index:

l[:index].index(element)

To find the nth index:

list(filter(lambda x:l[x]==element,range(len(l))))[n]

Answer 4

Use a list if you have multiple choices based on int

Say for example you have some output that will be 1, 0, or -1 and you need a different output for each case. You could do something like this:

print('0'if x==0else('1'if x>0else'-1'))

However, the better way is to use x as an index to a list like so:

print(['0','1','-1'][x])

which is 16 bytes shorter.

Answer 5

You can assign to a list inside of a for loop.

For example:

L=[1, 2, 3, 4, 6]
queue = [None]*len(L)
for e, queue[e] in enumerate(L):
    print("adding", queue[e], "to processing queue")

This can also be helpful if you need to switch the object you're assigning to.

class Foo:
    def __init__(self):
        self.x = None
a = Foo()
b = Foo()
for q, (lambda x: a if x%2==0 else b)(q).x in enumerate(range(10)):
    print(a.x, b.x)

Answer 6

Rename everything

Here's a little Python 2 snippet that takes a module and a string, and renames every function in that module whose name is longer than 2 characters to a single character with the provided string prefixed. If you're writing a VERY LONG python program that uses many library or builtin functions (and if you manage to golf this snippet better than I have), it has the potential to save quite a few characters. On short programs or programs that use few functions, it will be useless. Since dir() sorts the names in a module, this will always provide the same names to the same functions, and you can use globals() to inspect which names it has given to which functions.

import string
def _(x,y):
 for c,f in zip(string.letters,[x.__dict__[q]for q in dir(x)if q in x.__dict__ and(len(q)>2)*type(x.__dict__[q]).__name__.find('eth')>0]):globals()[y+c]=f

You can use it to rename all the string and builtin functions like so:

_(str,'s')
_(__builtins__,'')

And then see what you actually ended up naming them like so:

for k in sorted(globals().keys(),key=lambda x:`len(x)`+x):print k,globals()[k]

If you only want to rename the builtin functions, it's best not to define the function and just use the body directly:

import string
b=__builtins__
for c,f in zip(string.letters,[b.__dict__[q]for q in dir(b)if(len(q)>2)*type(x.__dict__[q]).__name__.find('eth')>0]):globals()[c]=f

Answer 7

Pathlib from shorter files manipulations:

# get current dir, go up one dir, go down one dir, list all "py" files
# and get the whole file bytes

import os.path as p
import glob

d = p.join(p.dirname(p.abspath(__file__))), 'foo', '*.py')
for x in glob.glob(d):
    with open(x, 'rb') as f:
        do_stuff(f)

Becomes:

import pathlib as p

d = p.Path(__file__).absolute().parent / 'foo'
for f in d.glob('*.py'):
    do_stuff(f.read_bytes())

Answer 8

Checking the length of a list

Empty      : `a==[]` but just checking if it's non empty and swapping the if and the else can be shorter 
Non-Empty  : `a` (assuming it is in a situation where it will be interpreted as a boolean)
len(a)>i   : `a>a[:i]` if the list is non-empty

Answer 9

Fuse list comprehensions

You can write

[f(x)for x in l]+[g(x)for x in l]

As

sum([[f(x),g(x)]for x in l],[])

It gets even better with more comprehensions or if you have to take out more values

If you need to expand a list you can even turn l+[f(x)for x in l]+[g(x)for x in l] into sum([[f(x),g(x)]for x in l],l)

Answer 10

1. from lib import func as F

2. from lib import*;F=func

3. import lib;F=lib.func

#2 is better than #1 except in rare cases where something in lib clobbers another name that's important to you.

#3 uses lib twice, winning with short library names.

Answer 11

I discovered a clever trick used here.

Instead of using the for loop to repeat multiple times, repeat exec multiple times.

p='+'
i=1
exec"print[p*i,i/9*p+'[>'+p*9+'<-]>'+i%9*p][i>20];i+=1;"*255

Compare this with

print"\n".join(">"+"+"*(i/16)+"[<"+"+"*16+">-]<"+"+"*(i%16)if i>31 else"+"*i for i in range(256))

Answer 12

Checking if a boolean value is false and an integer is a specific odd integer

~True is ~1 is -2 and ~False is ~0 is -1

-1&int is always int

-2&int is always even

Therefore, despite operator order, ~bool&int==odd_int works fine.

It's shorter than all of these equivalents:

not bool and int==odd_int

-~bool and int==odd_int

-~bool&int==odd_int

Answer 13

Do a ternary operator when you can.

Factorial example:

def f(x):return 1 if x==1 else x*f(x-1)

Answer 14

(Destructively) Iterate any iterable in reverse

for x in s[::-1]:f(x)

can shorten to:

while s:f(s.pop())

when s is a list and destructive iteration is acceptable. But if you need to use the value multiple times, even with the walrus, this doesn't work so hot:

for x in s[::-1]:f(x);g(x)

while s:x=s.pop();f(x);g(x)  # Longer!
while s:f(x:=s.pop());g(x)   # Same length, and := precedence means it doesn't always work in larger expressions

And of course, pop calls don't work at all when s isn't already a list. In Python 3 though, you can use unpacking to do the pop work with fewer characters and it works on things that aren't lists (though they become a list after the first loop):

for x in s[::-1]:f(x);g(x)
while s:*s,x=s;f(x);g(x)

which shaves two characters, and (since it's already doing top-level statement work) doesn't have the issues with precedence in larger expressions that the walrus has.

Even for the single expression case, this can be useful:

for x in s[::-1]:f(x) # Works on any sequence
while s:f(s.pop())    # 3 characters shorter but works only on lists
while s:*s,x=s;f(x)   # Only 2 characters shorter, but works on *any* iterable

Since the minimal number of characters to convert iterables to a list is 2-3 (depending on use), paying one extra character to avoid the conversion can be worth it.

Answer 15

Integer ceiling division with no repeated terms and no extra parentheses fighting operator precedence

You have n, the numerator, and d, the denominator, or an expression, a+b as the denominator. Floor division is trivial:

n//d      # 4 characters. On Py2, n/d works, but Py2 is dead, move on
n//(a+b)  # 8 characters if denominator is expression with operator that is lower precedence
          # than division (or equal and same associativity), dumb parens

The naive approaches to ceiling division is one of:

math.ceil(n/d)      # 14 characters, and needs an import. Wrong if inputs large enough to hit
                    # C double integer representation limits (~53 bits)
math.ceil(n/(a+b))  # 18 characters; bad, but at least the denominator being an expression
                    # didn't make it proportionately worse

(n+d-1)//d          # 10 characters. No import needed, no C double precision limitations
                    # Repeats d twice though, so if d is a more complicated expression...
(n+a+b-1)//(a+b)    # 16 characters, and gets worse if d is more complicated or involves
                    # operators with precedence below +/-

As mentioned off-hand in another answer, you can turn floor division into ceiling division with:

-(-n//d)        # 8 characters, shortest so far, and it only uses each term once, so...
-(-n//(a+b))    # 12 characters (whew, only paying for d expression once again)

but there is still one more optimization to be had. Unary negation is too high precedence, so we needed a second set of parentheses to separate negation of the numerator from the negation of the result of floor division (getting us ceil division as an end result). That's too much. But thankfully, subtraction is lower precedence than unary negation, so we can just subtract from zero, and get:

0--n//d      # 7 characters! Only three more than floor division!
0--n//(a+b)  # 11 characters, still only use d term once!

This is probably one of the only time Python code golfers are glad that Python doesn't offer a -- decrement operator; since it doesn't exist, each - is separate, the first is a subtraction operand (relatively low precedence), the second unary negation (very high precedence), so no parentheses are needed; unary negation attaches before floor division, subtraction after, achieving ceiling division at the cost of one 0, instead of open and close parentheses.

Answer 16

Use Annotations to save on indentation

Hurry up before this feature goes away

If you have code like this:

a:b = c

Is syntactically equivalent to

a = c
__annotations["a"] = b

Note that if b is a expression, it can be evaluated. You can use this for side effects, for example consider this python quine:

c:print(c%c)="c:print(c%%c)=%s"

In this case, this is the exact same length as

c="c:print(c%%c)=%s"
print(c%%c)

However, in a indented context, it can save some bytes:

if abcd:
   for defg in hijk:k:print(k(k))=eval(input())

Saves 5 bytes over:

if abcd:
    for defg in hijk:
        k=eval(input())
        print(k(k))

Answer 17

1 or 0 can act as boolean operators in Python:

func = lambda x:1 if x//2==x/2 else 0
while 1:
    if func(n):
         print('Hello')
    else:
         exit()

Which is 10 characters shorter than:

func = lambda x:True if x//2==x/2 else False
while True:
    if func(n):
         print('Hello')
    else:
         exit()

Answer 18

Not read all the answers but you can instead of

if x==3:
    print "yes"
else:
    print "no"

use

print "yes" if x==3 else "no"

Answer 19

Convert modules to lists

This will work for CPython (probably both 2 and 3), and lets you maybe shave a few bytes if you need to use a lot of different functions and classes with long names from the same module, but you aren't using any of them often enough to rename individually. You'll have to do some research first to figure out which magic numbers give you which functions. Example (rot13):

d=sorted
e=".__dict__.values()"
b=d(eval("__builtins__"+e))
s=d(eval("str"+e))
t=b[12]('string').maketrans
r=''.join(map(chr,range(65,91)))
w=r[13:]+r[:13]
l=t(r+s[4](r),w+s[4](w))
print s[28](b[53](),l)

Translated back to plain python, this is the same as:

t=__import__('string').maketrans
r=''.join(map(chr,range(65,91)))
w=r[13:]+r[:13]
l=t(r+r.lower(),w+w.lower())
print raw_input().translate(l)

which is obviously much shorter, but it should be clear how this methodology would eventually save bytes on much longer, more complicated programs that use more modules.

Answer 20

Slicing can assign:

Converting BGR pixels array to RGB:

l = len(pixels)
for idx in range(0, l - 2, 3):
    pixels[idx + 2], pixels[idx] = pixels[idx], pixels[idx + 2]

Becomes:

l = len(pixels)
pixels[2:l:3], pixels[0:l:3] = pixels[0:l:3], pixels[2:l:3]

It's not just shorter. It's 5 times faster. Except on pypy, where the first version is 2x times faster :)

Answer 21

Tricks with dicts:

# create a dict from iterable
d = {k: None for k in iterable}
# merge 2 dicts
d.merge(d1)
d.merge(d2)
# access a key, if it doesn't exist set a default value and return it
if 'foo' not in d:
    res = d['foo'] = 'bar'
    res = 'bar
else:
    res = d['foo']

Becomes:

d = {**dict.fromkeys(iterable), **d1, **d2}
res = d.setdefault('foo', 'bar')

Or if you need repeated access:

import collections as c
d = c.ChainMap(dict.fromkeys(iterable), d1, d2, {'foo': 'bar'})
res = d['foo']

Answer 22

When you want to compare 2 booleans to see if they are different, it's preferable to use ^ (XOR) rather than != (it can save one byte)

Example with: a: bool and b: bool

if a!=b: # 8 bytes

if a^b: # 7 bytes

Answer 23

Use **.5 instead of math.sqrt()

This is a very useful tip, especially if square roots are the base of your answer.

Instead of

from math import*
x=sqrt(4)

which is 27 bytes, we can simply use

x=4**0.5

Golfed, of course, it's

x=4**.5

saving 20 bytes!

Note that this trick can be used for finding any root. Here is a function that does this:

lambda x,y:x**1/y

Answer 24

Another tip: this is a pretty useful (but rarely useable) alternative ternary operator which we use in Scratch: * and +. This uses the fact the booleans are equal to 0 and 1, and the rules of multiplying, but can only be used in limited situations.

For example, in string manipulations, you can use these. I used this in my code.golf submission:

for i in range(1,101):print(f'{i}'*bool(i%3and i%5)+'Fizz'*(i%3==0)+'Buzz'*(i%5==0))

These can be sometimes long, so you can only use in limited situations. Note that string conversions may be necessary, as used in f'{i}', thus making this impractical with multiple datatypes.

This is also dangerous in recursive solutions, as all the conditions will be checked (so it doesn't terminate at a True value, thus recursing unnecessarily; the exact opposite of a short-circuit operator).

So in factorials, this will fail to RecursionError:

x=lambda n:n*(x(n-1)*(n!=1)+1*(n==1))

Answer 25

Should you set variable?

_{This is perhaps already be posted but I could not find it with an explanation.}

gsetn = lambda n, k, c=3: n+c < (n-1) * k
gset = lambda n, s, c=3: gsetn(n, len(s), c)

gset(3,'range')       # == True
gsetn(3, 4)           # == True

The cost of using f of length k n times is nk and if it is set, the cost is n+cost of setting. So if f is used n times it is worth setting if $$ n+c+k < n\cdot k\Leftrightarrow n+c < (n-1)\cdot k. $$ Typically it costs x=+k+\n(or ;) = 3 + k to set the variable. Example:

len('r,a,s=range,abs,sum;') == len('range')+len('abs')+len('sum')+3*3

But, it could be different:

from math import factorial as f

Answer 26

From what i learned:

1. Encode your long codes to base64 and use eval to run them: eval(ENCODEFUNC(CODE)). Something like zlib has half the size of b64 but there are even more efficient hash function that you can look up.
2. Use ternary operators and lambdas.
3. If the the task you trying to accomplish is a trivial thing then probably there is a builtin function for it somewhere.
4. Always look for a case for using list comprehensions and generators stuff.