weighted_distribution.py

import unittest
from typing import List, Dict
from tqdm import tqdm

def distribute_elements(stream: List[int], servers: List[str], weights: List[int]) -> Dict[str, List[int]]:
    if len(servers) != len(weights):
        raise ValueError("The number of servers and weights must be the same.")
    
    total_weight = sum(weights)
    normalized_weights = [weight / total_weight for weight in weights]
    
    distribution = {server: [] for server in servers}
    
    for element in tqdm(stream, desc='Distributing elements'):
        min_server = min(distribution, key=lambda x: (len(distribution[x]) / normalized_weights[servers.index(x)] if normalized_weights[servers.index(x)] > 0 else float('inf')))
        distribution[min_server].append(element)
    
    return distribution

# Unit test for the function
class TestDistribution(unittest.TestCase):
    def test_distribute_elements(self):
        stream = range(100)
        servers = ['Server1', 'Server2', 'Server3']
        weights = [5, 3, 2]
        expected_distribution = {
            'Server1': [0, 3, 5, 7, 9, 10, 13, 15, 17, 19, 20, 23, 25, 27, 29, 30, 33, 35, 37, 39, 40, 43, 45, 47, 49, 50, 53, 55, 57, 59, 60, 63, 65, 67, 69, 70, 73, 75, 77, 79, 80, 83, 85, 87, 89, 90, 93, 95, 97, 99],
            'Server2': [1, 4, 8, 11, 14, 18, 21, 24, 28, 31, 34, 38, 41, 44, 48, 51, 54, 58, 61, 64, 68, 71, 74, 78, 81, 84, 88, 91, 94, 98],
            'Server3': [2, 6, 12, 16, 22, 26, 32, 36, 42, 46, 52, 56, 62, 66, 72, 76, 82, 86, 92, 96]
        }
        result_distribution = distribute_elements(stream, servers, weights)
        self.assertEqual(result_distribution, expected_distribution)

# Run the unit test without the progress bar (for testing purposes)
unittest.main(argv=[''], verbosity=2, exit=False)

weighted_random.py

import random
import unittest
from typing import List, Dict

def distribute_elements_randomly(stream: List[int], servers: List[str], weights: List[int]) -> Dict[str, List[int]]:
    if len(servers) != len(weights):
        raise ValueError("The number of servers and weights must be the same.")
    
    # Normalize weights
    total_weight = sum(weights)
    normalized_weights = [weight / total_weight for weight in weights]
    
    # Initialize the distribution
    distribution = {server: [] for server in servers}
    
    for element in stream:
        # Choose a server based on weights
        chosen_server = random.choices(servers, weights=normalized_weights)[0]
        distribution[chosen_server].append(element)
    
    return distribution

# Unit test for the function with randomization
class TestRandomDistribution(unittest.TestCase):
    def test_distribute_elements_randomly(self):
        stream = list(range(100))
        servers = ['Server1', 'Server2', 'Server3']
        weights = [5, 3, 2]
        distribution = distribute_elements_randomly(stream, servers, weights)
        
        # Check that every element is distributed
        distributed_elements = [item for sublist in distribution.values() for item in sublist]
        self.assertCountEqual(stream, distributed_elements)
        
        # Check that the distribution respects the weight ratio within a margin
        total_elements = len(stream)
        expected_ratios = [weight / sum(weights) for weight in weights]
        actual_ratios = [len(distribution[server]) / total_elements for server in servers]
        
        for expected, actual in zip(expected_ratios, actual_ratios):
            self.assertAlmostEqual(expected, actual, delta=0.1)

# Run the unit tests
unittest.main(argv=[''], exit=False)