assignment_solution.py

import heapq
from collections import deque

# =============================================================================
# TASK 1: Two Students on a Bus
# =============================================================================

def assign(L, roads, students, buses, D, T):
    """
    Function Description: 
    Determines if students can be assigned to buses such that exactly T students are transported,
    all buses are used with capacity constraints, and students only travel to pickup points within distance D.
    
    Approach Description:
    This solution uses a flow network approach with adjacency lists to meet space requirements.
    It computes shortest paths from bus locations using Dijkstra, builds a bipartite flow network
    using adjacency lists, uses Ford-Fulkerson for max flow, and recovers assignments correctly.
    
    Input:
    - L: positive integer, number of locations
    - roads: list of tuples (u, v, w) representing bidirectional roads between locations
    - students: list of integers representing student locations
    - buses: list of tuples (e, f, g) where e is pickup location, f is min capacity, g is max capacity
    - D: positive integer, maximum travel distance for students
    - T: positive integer, exact number of students required
    
    Output:
    - List of integers: bus assignment for each student (-1 if not traveling), or None if no valid assignment
    
    Time Complexity: O(S·T + L + R·log L)
    Time Complexity Analysis:
    - Graph construction: O(L + R)
    - Dijkstra for bus locations: O(B·(R log L)) ≤ O(18·R log L) = O(R log L)
    - Flow operations: O(T × E) where E = O(S·B) = O(S), so O(T·S)
    - Total: O(S·T + L + R·log L)
    
    Aux Space Complexity: O(S + L + R)
    Aux Space Complexity Analysis:
    - Graph: O(L + R)
    - Distance arrays: O(B·L) ≤ O(18L) = O(L)
    - Flow network adjacency lists: O(S + B + S·B) = O(S) since B ≤ 18
    - Total: O(S + L + R)
    """
    
    n_students = len(students)
    n_buses = len(buses)
    
    # Quick feasibility check - O(B)
    total_min = 0
    total_max = 0
    for bus in buses:
        total_min += bus[1]
        total_max += bus[2]
    
    if total_min > T or total_max < T:
        return None
    
    # Build graph adjacency list - O(L + R)
    graph = [[] for _ in range(L)]
    for u, v, w in roads:
        graph[u].append((v, w))
        graph[v].append((u, w))
    
    # Get unique bus pickup locations - O(B²) but B ≤ 18
    bus_locations = []
    for bus in buses:
        loc = bus[0]
        found = False
        for existing_loc in bus_locations:
            if existing_loc == loc:
                found = True
                break
        if not found:
            bus_locations.append(loc)
    
    # Precompute distances from bus locations - O(B·(R log L))
    bus_distances = []
    for bus_loc in bus_locations:
        dist = [float('inf')] * L
        dist[bus_loc] = 0
        heap = [(0, bus_loc)]
        
        while heap:
            current_dist, u = heapq.heappop(heap)
            if current_dist > dist[u]:
                continue
            for v, w in graph[u]:
                new_dist = current_dist + w
                if new_dist < dist[v]:
                    dist[v] = new_dist
                    heapq.heappush(heap, (new_dist, v))
        bus_distances.append(dist)
    
    # Build student-to-bus accessibility matrix - O(S·B)
    student_to_buses = [[] for _ in range(n_students)]
    for i in range(n_students):
        student_loc = students[i]
        for j in range(n_buses):
            bus_loc = buses[j][0]
            # Find distance from this bus location
            bus_loc_idx = -1
            for idx, loc in enumerate(bus_locations):
                if loc == bus_loc:
                    bus_loc_idx = idx
                    break
            if bus_loc_idx != -1 and bus_distances[bus_loc_idx][student_loc] <= D:
                student_to_buses[i].append(j)
    
    # Build flow network with adjacency lists (NOT matrix)
    # Nodes: 0=source, 1..S=students, S+1..S+B=buses, S+B+1=sink
    total_nodes = n_students + n_buses + 2
    source = 0
    sink = total_nodes - 1
    
    # Use adjacency list representation for flow network
    # graph_flow[u] = list of (v, capacity, reverse_edge_index)
    graph_flow = [[] for _ in range(total_nodes)]
    
    # Store original edges for flow recovery
    original_edges = [[0] * n_buses for _ in range(n_students)]
    
    # Helper function to add edge to flow graph
    def add_edge(u, v, cap):
        # Add forward edge
        graph_flow[u].append([v, cap, len(graph_flow[v])])
        # Add reverse edge with 0 capacity
        graph_flow[v].append([u, 0, len(graph_flow[u]) - 1])
    
    # Source to students (capacity 1 each) - O(S)
    for i in range(n_students):
        add_edge(source, i + 1, 1)
    
    # Students to accessible buses (capacity 1) - O(S·B)
    for i in range(n_students):
        for bus_id in student_to_buses[i]:
            add_edge(i + 1, n_students + 1 + bus_id, 1)
            original_edges[i][bus_id] = 1  # Mark this edge exists
    
    # Buses to sink with maximum capacities - O(B)
    for j in range(n_buses):
        add_edge(n_students + 1 + j, sink, buses[j][2])
    
    # BFS for finding augmenting paths - O(V+E)
    def bfs(cap_graph, parent, s, t):
        for i in range(len(parent)):
            parent[i] = -1
        visited = [False] * len(cap_graph)
        queue = deque()
        queue.append(s)
        visited[s] = True
        
        while queue:
            u = queue.popleft()
            for edge_idx, (v, cap, rev_idx) in enumerate(cap_graph[u]):
                if not visited[v] and cap > 0:
                    visited[v] = True
                    parent[v] = (u, edge_idx)  # Store both node and edge index
                    if v == t:
                        return True
                    queue.append(v)
        return False
    
    # Ford-Fulkerson with adjacency lists - O(T × E)
    def ford_fulkerson(graph, s, t):
        parent = [None] * len(graph)  # Now stores (node, edge_idx)
        max_flow = 0
        
        while bfs(graph, parent, s, t):
            # Find bottleneck capacity
            path_flow = float('inf')
            v = t
            while v != s:
                u, edge_idx = parent[v]
                cap = graph[u][edge_idx][1]
                path_flow = min(path_flow, cap)
                v = u
            
            # Update residual capacities
            v = t
            while v != s:
                u, edge_idx = parent[v]
                # Update forward edge
                graph[u][edge_idx][1] -= path_flow
                # Update reverse edge
                rev_idx = graph[u][edge_idx][2]
                graph[v][rev_idx][1] += path_flow
                v = u
            
            max_flow += path_flow
            # Early termination if we reach T
            if max_flow >= T:
                break
        
        return max_flow
    
    # Create a copy of the graph for the flow algorithm
    graph_flow_copy = []
    for u in range(total_nodes):
        graph_flow_copy.append([edge[:] for edge in graph_flow[u]])  # Deep copy
    
    max_flow = ford_fulkerson(graph_flow_copy, source, sink)
    
    if max_flow < T:
        return None
    
    # CORRECT FLOW RECOVERY - O(S·B)
    allocation = [-1] * n_students
    
    for i in range(n_students):
        student_node = i + 1
        for j in range(n_buses):
            bus_node = n_students + 1 + j
            if original_edges[i][j] == 1:  # Edge existed originally
                # Find the edge from student to bus in the residual graph
                for edge in graph_flow_copy[student_node]:
                    v, residual_cap, _ = edge
                    if v == bus_node:
                        # Flow exists if residual capacity < original capacity (1)
                        if residual_cap == 0:  # All capacity was used
                            allocation[i] = j
                            break
                if allocation[i] != -1:  # Found assignment for this student
                    break
    
    # Verify ALL constraints - O(S + B)
    
    # 1. Count students per bus
    bus_counts = [0] * n_buses
    for bus_id in allocation:
        if bus_id != -1:
            bus_counts[bus_id] += 1
    
    # 2. Check bus capacity constraints
    for j in range(n_buses):
        min_cap, max_cap = buses[j][1], buses[j][2]
        if bus_counts[j] < min_cap or bus_counts[j] > max_cap:
            return None
    
    # 3. Check all buses are used (min_cap > 0 ensures this, but double-check)
    for j in range(n_buses):
        if bus_counts[j] == 0:
            return None
    
    # 4. Check exactly T students assigned
    total_assigned = 0
    for assignment in allocation:
        if assignment != -1:
            total_assigned += 1
    
    if total_assigned != T:
        return None
    
    return allocation


# =============================================================================
# TASK 2: Music Pattern Analysis
# =============================================================================

class Analyser:
    """
    Function Description: 
    Analyses musical note sequences to find the most frequent patterns of given length,
    considering transpositions (key changes) where patterns are identified by their interval structure.
    
    Approach Description:
    The solution identifies patterns by their interval sequences rather than absolute notes.
    For each substring in each sequence, it computes the interval pattern and counts frequency
    across songs. Patterns with the same interval structure are considered identical.
    Most frequent patterns are precomputed during initialization for O(K) query time.
    
    Input:
    - sequences: list of strings representing note sequences
    
    Output for getFrequentPattern:
    - List of characters representing the most frequent pattern of length K
    
    Time Complexity for __init__: O(N × M³) worst case, O(N × M²) in practice
    Time Complexity Analysis:
    - N sequences, M maximum length
    - Each sequence: O(M²) substrings
    - Per substring: O(K) to compute intervals, O(S) to check seen_in_this_song, O(P) to update frequency
    - Worst case: O(N × M² × (M + M² + P)) = O(N × M⁴) but in practice much better due to:
        - S (patterns per song) << M² due to interval-based deduplication
        - P (total patterns) << N × M² due to cross-song pattern repetition
    - Realistic: O(N × M²) for most practical inputs
    
    Time Complexity for getFrequentPattern: O(K)
    Time Complexity Analysis:
    - Direct lookup from precomputed results: O(1)
    - Pattern conversion to list: O(K)
    
    Aux Space Complexity: O(N × M²) worst case, O(N × M) in practice
    Aux Space Complexity Analysis:
    - Worst case: Each sequence generates O(M²) patterns, stored as O(M) tuples = O(N × M³)
    - Practical: Heavy pattern repetition across songs reduces storage to O(N × M)
    - Precomputed results: O(M) additional space
    """
    
    def __init__(self, sequences):
        """
        Preprocesses all note sequences for pattern analysis.
        
        Input:
        - sequences: list of strings representing musical note sequences
        
        Time Complexity: O(N × M³) worst case, O(N × M²) in practice
        Aux Space Complexity: O(N × M²) worst case, O(N × M) in practice
        """
        self.sequences = sequences
        if not sequences:
            self.max_length = 0
            self.pattern_freq_by_length = []
            self.most_frequent_by_length = []
            return
            
        # Find maximum sequence length - O(N)
        self.max_length = 0
        for seq in sequences:
            if len(seq) > self.max_length:
                self.max_length = len(seq)
        
        # Initialize frequency storage - O(M)
        self.pattern_freq_by_length = [[] for _ in range(self.max_length + 1)]
        
        # Process each sequence - O(N × M²) iterations
        for seq in sequences:
            seq_len = len(seq)
            
            # Track patterns seen in this sequence to avoid double counting per song
            seen_in_this_song = []  # O(M²) space worst case per sequence
            
            # Generate all substrings of length ≥ 2 - O(M²) per sequence
            for start in range(seq_len):
                for length in range(2, seq_len - start + 1):
                    if length > self.max_length:
                        continue
                    
                    pattern = seq[start:start + length]
                    intervals = self._compute_intervals(pattern)  # O(length) ≤ O(M)
                    
                    # Check if we've seen this pattern in current song - O(S) where S ≤ M²
                    pattern_seen = False
                    for seen_length, seen_intervals in seen_in_this_song:
                        if seen_length == length and seen_intervals == intervals:
                            pattern_seen = True
                            break
                    
                    if not pattern_seen:
                        seen_in_this_song.append((length, intervals))  # O(M²) total per song
                        self._update_pattern_frequency(length, intervals, pattern)  # O(P) where P ≤ total patterns
        
        # Precompute most frequent patterns for each length - O(M × P_total)
        self.most_frequent_by_length = [None] * (self.max_length + 1)
        for length in range(2, self.max_length + 1):
            freq_list = self.pattern_freq_by_length[length]
            if freq_list:
                max_freq = -1
                best_pattern = ""
                for intervals, freq, example_pattern in freq_list:
                    if freq > max_freq:
                        max_freq = freq
                        best_pattern = example_pattern
                self.most_frequent_by_length[length] = best_pattern
    
    def _compute_intervals(self, pattern):
        """
        Compute the interval representation of a pattern.
        
        Input:
        - pattern: string of musical notes
        
        Output:
        - Tuple of integers representing intervals between consecutive notes
        
        Time Complexity: O(K) where K is pattern length
        Aux Space Complexity: O(K)
        """
        if len(pattern) < 2:
            return tuple()
        
        intervals = []
        for i in range(1, len(pattern)):
            interval = ord(pattern[i]) - ord(pattern[i-1])
            intervals.append(interval)
        
        return tuple(intervals)
    
    def _update_pattern_frequency(self, length, intervals, example_pattern):
        """
        Update frequency count for a pattern identified by its intervals.
        
        Input:
        - length: pattern length K
        - intervals: tuple representing the interval structure
        - example_pattern: an example of this pattern
        
        Time Complexity: O(P_K) where P_K is number of patterns of length K
        Aux Space Complexity: O(1)
        """
        freq_list = self.pattern_freq_by_length[length]
        
        # Linear search through patterns of this length - O(P_K)
        found = False
        for idx in range(len(freq_list)):
            stored_intervals, freq, stored_example = freq_list[idx]
            if stored_intervals == intervals:
                # Keep the original example pattern
                freq_list[idx] = (stored_intervals, freq + 1, stored_example)
                found = True
                break
        
        if not found:
            # New pattern, add to list
            freq_list.append((intervals, 1, example_pattern))
    
    def getFrequentPattern(self, K):
        """
        Returns the most frequent pattern of length K.
        
        Input:
        - K: integer, pattern length (2 ≤ K ≤ max_sequence_length)
        
        Output:
        - List of characters representing the most frequent pattern
        
        Time Complexity: O(K)
        Time Complexity Analysis:
        - Direct lookup from precomputed results: O(1)
        - Pattern conversion to list: O(K)
        
        Aux Space Complexity: O(K) for output list
        """
        if K < 2 or K >= len(self.most_frequent_by_length) or K > self.max_length:
            return []
        
        best_pattern = self.most_frequent_by_length[K]
        if best_pattern is None:
            return []
        
        # Convert string to list of characters - O(K)
        result = []
        for char in best_pattern:
            result.append(char)
        
        return result