Full implementation of the mover function

How the mover function is implemented is beyond the scope of the mathematics of this book, but the implementation is provided here (Note that this currently only works over two or three inputs, and it is future work to extend this to work over more qubits.):

def diffusion_step(qr,qc,num_inputs):
    if num_inputs not in [2,3]:
        raise Exception("currently only supports 2 or 3 inputs")
    for i in range(num_inputs):
        qc.h(qr[i])
    for i in range(num_inputs): 
        qc.x(qr[i])
 
    control_Z(qr,qc,num_inputs)

    for i in range(num_inputs):
        qc.x(qr[i])
    for i in range(num_inputs):
        qc.h(qr[i])

The mover function relies on a controlled Z gate, which is implemented as follows:

def control_Z(qr,qc,num_inputs):
    if num_inputs not in [2,3]:
        raise Exception("currently only supports 2 or 3 inputs")
    if num_inputs==2:
        qc.h(qr[1])
        qc.cx(qr[0],qr[1])
        qc.h(qr[1])
    elif num_inputs==3:
        qc.h(qr[2])
        qc.ccx(qr[0],qr[1],qr[2]) 
        qc.h(qr[2])