EnjunDu/GraphMaster

1import json
2
3def clean_graph_data(input_file, output_file):
4    """
5    Clean graph data by removing nodes with mask="None", ensuring sequential node_ids,
6    and updating all neighbor references accordingly.
7    """
8    # Load the JSON data
9    with open(input_file, 'r', encoding='utf-8') as f:
10        data = json.load(f)
11    
12    # Identify valid nodes (mask != "None") and their IDs
13    valid_nodes = []
14    valid_node_ids = set()
15    
16    for node in data:
17        if 'mask' in node and node['mask'] != "None":
18            valid_nodes.append(node)
19            valid_node_ids.add(node['node_id'])
20    
21    # Create mapping from old node_id to new node_id
22    old_to_new_mapping = {}
23    new_id = 0
24    
25    for node in sorted(valid_nodes, key=lambda x: x['node_id']):
26        old_to_new_mapping[node['node_id']] = new_id
27        new_id += 1
28    
29    # Update node_ids and neighbors based on the mapping
30    for node in valid_nodes:
31        # Update neighbors first (while node_id is still the old one)
32        new_neighbors = []
33        for neighbor in node['neighbors']:
34            if neighbor in valid_node_ids:  # Only keep neighbors that weren't removed
35                new_neighbors.append(old_to_new_mapping[neighbor])
36        node['neighbors'] = new_neighbors
37        
38        # Update node_id
39        node['node_id'] = old_to_new_mapping[node['node_id']]
40    
41    # Sort nodes by new node_id for better readability
42    valid_nodes.sort(key=lambda x: x['node_id'])
43    
44    # Save the cleaned data
45    with open(output_file, 'w') as f:
46        json.dump(valid_nodes, f, indent=2)
47    
48    return f"Successfully cleaned the graph data. Removed {len(data) - len(valid_nodes)} nodes with mask='None'."
49
50# Usage
51result = clean_graph_data('wikics.json', 'wikics_cleaned.json')
52print(result)