try to improve two_sum.py