babylon_tower.py

import time

# Square Pyramidal Numbers
# Problem: https://github.com/GeekBrainsTutorial/Python_lessons_basic/blob/master/lesson02/home_work/hw02_hard.py


def get_room_v1(number):
    last_room = 0
    section = 0
    # last_floor = 0
    all_floors = list()
    all_rooms = list()
    while last_room <= number:
        section += 1
        rooms_in_section = section ** 2
        all_floors.append(section)
        all_rooms.append(rooms_in_section)
        last_room = sum(all_rooms)
        # last_floor = sum(all_floors)
    # print section, last_floor, last_room
    first_in_section = last_room - (section ** 2 - 1)
    floors_in_section = list()
    for floor in xrange(section):
        floors_in_section.append(range(first_in_section, first_in_section + section))
        first_in_section += section
    for floor in floors_in_section:
        # print floor
        if number in floor:
            target_floor = section + floors_in_section.index(floor) + 1
            position = floor.index(number) + 1
            return target_floor, position


def binary_search(num, first, last, dim):
    middle = round((first + last) / 2.0)
    # print 'middle', middle
    # print 'first', first
    # print 'last', last
    last_on_floor = middle * dim
    # print last_on_floor
    first_on_floor = last_on_floor - dim + 1
    # print first_on_floor
    if first_on_floor <= num <= last_on_floor:
        floor = int(middle)
        position = num - first_on_floor + 1
        return floor, int(position)
    else:
        # print 'num:', num
        if num < first_on_floor:
            # print 'num <', first_on_floor
            # print num, first, int(middle - 1)
            return binary_search(num, first, int(middle - 1), dim)
        elif num > last_on_floor:
            # print 'num >', last_on_floor
            return binary_search(num, middle + 1, last, dim)


def get_room_v2(number):
    sequence = list()
    section = 0
    last = 0
    while number >= last:
        section += 1
        last = (section * (section + 1) * (2 * section + 1)) / 6
        sequence.append(last)
    # last_room = sequence[-1]
    first_room = sequence[-2] + 1
    last_floor = (section * (section + 1)) / 2
    first_floor = last_floor - section + 1
    # print section, first_room, last_room, first_floor, last_floor
    abs_target = number - first_room + 1
    # abs_first_room = 1
    # abs_last_room = section ** 2
    abs_first_floor = 1
    result = binary_search(abs_target, abs_first_floor, section, section)
    floor = result[0] + first_floor - 1
    return floor, result[1]


def binary_search_iter(num, first, last):
    bottom = first
    top = last
    # print num
    while True:
        middle = int(round((bottom + top) / 2.0))
        l_room = (middle * (middle + 1) * (2 * middle + 1)) / 6
        f_room = l_room - middle ** 2 + 1
        # print middle, f_room, l_room
        if f_room <= num <= l_room:
            return middle, f_room, l_room
        if num > l_room:
            bottom = middle + 1
        elif num < f_room:
            top = middle - 1


def get_room_v3(number):
    section = 0
    last_room = 0
    while number >= last_room:
        section += 1
        last_room = (section * (section + 1) * (2 * section + 1)) / 6
    first_room = last_room - section ** 2 + 1
    last_floor = (section * (section + 1)) / 2
    first_floor = last_floor - section + 1
    # print section, first_room, last_room, first_floor, last_floor
    abs_target = number - first_room + 1
    abs_first_floor = 1
    result = binary_search(abs_target, abs_first_floor, section, section)
    floor = result[0] + first_floor - 1
    return floor, result[1]


def get_room_v4(num):
    section_rooms = binary_search_iter(num, 1, num)
    section = section_rooms[0]
    first_room = section_rooms[1]
    last_floor = (section * (section + 1)) / 2
    first_floor = last_floor - section + 1
    abs_target = num - first_room + 1
    abs_first_floor = 1
    result = binary_search(abs_target, abs_first_floor, section, section)
    floor = result[0] + first_floor - 1
    return floor, result[1]


if __name__ == '__main__':
    # print get_room(3)
    start = time.time()
    # print get_room_v2(10 ** 20)
    # print get_room_v3(10 ** 20)
    print get_room_v4(10 ** 30)
    stop = time.time()
    print 'running time:', stop - start
    # print binary_search(8, 1, 4, 4)
    # print binary_search_iter(10, 1, 10)