用户关系链

import random
import math


def get_follows(num2users, user, follow_num):
    follows = []
    last_users = []
    for i in list(reversed(list(range(len(num2users))))):
        keeps = []
        users = num2users[i]
        if user in users:
            keeps.append(user)
            users.remove(user)

        user_num = len(users)
        if follow_num > user_num:
            selects = users
            follow_num -= user_num
        else:
            selects = users[:follow_num]
            keeps += users[follow_num:]
            follow_num = 0

        keeps += last_users
        follows += selects
        last_users = selects
        num2users[i] = keeps
    while len(num2users[len(num2users) - 1]) == 0:
        num2users = num2users[:len(num2users) - 1]
    return num2users, follows


# 20个测试 # 缺30个
if __name__ == '__main__':
    user_num = 500000 - 50
    max_follow = 20

    user_num = 20
    max_follow = 19

    users = list(range(user_num))

    follow_nums = [random.randint(1, max_follow) for _ in range(len(users))]
    follow_nums = sorted(follow_nums, reverse=True)
    fans_nums = list(follow_nums)
    random.shuffle(fans_nums)

    infos = list(zip(users, follow_nums, fans_nums))

    num2users = [[] for i in range(max_follow + 1)]
    for user, _, fans_num in infos:
        num2users[fans_num].append(user)

    for i, (user, follow_num, _) in enumerate(infos):
        print(i + 1, len(infos))
        num2users, follows = get_follows(num2users, user, follow_num)
        if user in follows:
            follows.remove(user)
        infos[i] = list(infos[i])
        infos[i].append(follows)